diff --git a/CHANGELOG.md b/CHANGELOG.md index f2ffbc9fda..b54f164594 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -83,6 +83,12 @@ - Export `ViewerBase` from `newton.viewer` public API - Add `custom_attributes` argument to `ModelBuilder.add_shape_convex_hull()` - Add RJ45 plug-socket insertion example with SDF contacts, latch joint, and interactive gizmo +- Add `TRIANGLE_PRISM` support-function type for heightfield triangles, extruding 1 m along the heightfield's local -Z so GJK/MPR naturally resolves shapes on the back side +- Add `ViewerGL.log_scalar()` for live scalar time-series plots in the viewer +- Add `deterministic` flag to `CollisionPipeline` and `NarrowPhase` for GPU-thread-scheduling-independent contact ordering via radix sort and deterministic fingerprint tiebreaking in contact reduction +- Add `ViewerBase.log_arrows()` for arrow rendering (wide line + arrowhead) in the GL viewer with a dedicated geometry shader +- Add `SolverVBD.set_joint_constraint_mode()` and `SolverVBD.JointSlot` for runtime hard/soft constraint toggling per joint +- Add `rigid_contact_hard`, `rigid_contact_history`, `rigid_avbd_contact_alpha`, `rigid_avbd_joint_alpha` params to `SolverVBD` for augmented-Lagrangian contact and joint control ### Changed @@ -112,6 +118,11 @@ - Use `Literal` types for `SolverImplicitMPM.Config` string fields with fixed option sets (`solver`, `warmstart_mode`, `collider_velocity_mode`, `grid_type`, `transfer_scheme`, `integration_scheme`) - Migrate `wp.array(dtype=X)` type annotations to `wp.array[X]` bracket syntax (Warp 1.12+). - Align articulated `State.body_qd` / FK / IK / Jacobian / mass-matrix linear velocity with COM-referenced motion. If you were comparing `body_qd[:3]` against finite-differenced body-origin motion, recover origin velocity via `v_origin = v_com - omega x r_com_world`. Descendant `FREE` / `DISTANCE` `joint_qd` remains parent-frame and `joint_f` remains a world-frame COM wrench. +- Pin `mujoco` and `mujoco-warp` dependencies to `~=3.6.0` +- Update default environment map texture in GL viewer (source: https://polyhaven.com/a/brown_photostudio_02) +- Increase conveyor rail roughness in `example_basic_conveyor` to reduce mirror-like reflections +- Reduce default `stretch_stiffness` from `1.0e9` to `1.0e4` in `add_rod()` and `add_rod_graph()` +- VBD solver uses augmented-Lagrangian hard constraints for body-body contacts by default (`rigid_contact_hard=True`) ### Deprecated @@ -129,6 +140,7 @@ - Deprecate `ModelBuilder.add_shape_ellipsoid()` parameters `a`, `b`, `c` in favor of `rx`, `ry`, `rz` - Deprecate passing a `Gaussian` as the second positional argument to `ModelBuilder.add_shape_gaussian()`; use the `gaussian=` keyword argument instead - Deprecate `SensorTiledCamera.utils.assign_random_colors_per_world()` and `assign_random_colors_per_shape()` in favor of per-shape colors via `ModelBuilder.add_shape_*(color=...)` +- Deprecate `rigid_enable_dahl_friction` in `SolverVBD`; Dahl friction is now auto-detected from model attributes (`model.vbd.dahl_eps_max` / `model.vbd.dahl_tau`) ### Removed diff --git a/docs/images/examples/example_cable_pile.jpg b/docs/images/examples/example_cable_pile.jpg index 8c2f7f804e..566c1337a0 100644 Binary files a/docs/images/examples/example_cable_pile.jpg and b/docs/images/examples/example_cable_pile.jpg differ diff --git a/newton/_src/sim/builder.py b/newton/_src/sim/builder.py index 340240ab60..cfdddf4c21 100644 --- a/newton/_src/sim/builder.py +++ b/newton/_src/sim/builder.py @@ -4219,7 +4219,7 @@ def add_joint_cable( """ # Linear DOF (stretch) - se_ke = 1.0e9 if stretch_stiffness is None else stretch_stiffness + se_ke = 1.0e4 if stretch_stiffness is None else stretch_stiffness se_kd = 0.0 if stretch_damping is None else stretch_damping ax_lin = ModelBuilder.JointDofConfig(target_ke=se_ke, target_kd=se_kd) @@ -6427,7 +6427,7 @@ def add_rod( cfg = self.default_shape_cfg # Stretch defaults: high stiffness to keep neighboring capsules tightly coupled - stretch_stiffness = 1.0e9 if stretch_stiffness is None else stretch_stiffness + stretch_stiffness = 1.0e4 if stretch_stiffness is None else stretch_stiffness stretch_damping = 0.0 if stretch_damping is None else stretch_damping # Bend defaults: 0.0 (users must explicitly set for bending resistance) @@ -6605,7 +6605,7 @@ def add_rod_graph( cfg = self.default_shape_cfg # Stretch defaults: high stiffness to keep neighboring capsules tightly coupled - stretch_stiffness = 1.0e9 if stretch_stiffness is None else stretch_stiffness + stretch_stiffness = 1.0e4 if stretch_stiffness is None else stretch_stiffness stretch_damping = 0.0 if stretch_damping is None else stretch_damping # Bend defaults: 0.0 (users must explicitly set for bending resistance) diff --git a/newton/_src/solvers/vbd/particle_vbd_kernels.py b/newton/_src/solvers/vbd/particle_vbd_kernels.py index 7fec9c4ec3..4753cb24a4 100644 --- a/newton/_src/solvers/vbd/particle_vbd_kernels.py +++ b/newton/_src/solvers/vbd/particle_vbd_kernels.py @@ -16,7 +16,7 @@ import warp as wp from newton._src.math import orthonormal_basis -from newton._src.solvers.vbd.rigid_vbd_kernels import evaluate_body_particle_contact +from newton._src.solvers.vbd.rigid_vbd_kernels import _eval_body_particle_contact, evaluate_body_particle_contact from ...geometry import ParticleFlags from ...geometry.kernels import triangle_closest_point @@ -2439,7 +2439,7 @@ def accumulate_contact_force_and_hessian_no_self_contact( contact_kd = body_particle_contact_material_kd[t_id] contact_mu = body_particle_contact_material_mu[t_id] - body_contact_force, body_contact_hessian = evaluate_body_particle_contact( + body_contact_force, body_contact_hessian = _eval_body_particle_contact( particle_idx, pos[particle_idx], pos_anchor[particle_idx], @@ -2449,7 +2449,6 @@ def accumulate_contact_force_and_hessian_no_self_contact( contact_mu, friction_epsilon, particle_radius, - shape_material_mu, shape_body, body_q, body_q_prev, @@ -2940,7 +2939,7 @@ def accumulate_particle_body_contact_force_and_hessian( contact_kd = body_particle_contact_material_kd[t_id] contact_mu = body_particle_contact_material_mu[t_id] - body_contact_force, body_contact_hessian = evaluate_body_particle_contact( + body_contact_force, body_contact_hessian = _eval_body_particle_contact( particle_idx, pos[particle_idx], pos_anchor[particle_idx], @@ -2950,7 +2949,6 @@ def accumulate_particle_body_contact_force_and_hessian( contact_mu, friction_epsilon, particle_radius, - shape_material_mu, shape_body, body_q, body_q_prev, diff --git a/newton/_src/solvers/vbd/rigid_vbd_kernels.py b/newton/_src/solvers/vbd/rigid_vbd_kernels.py index 7f777d70bd..e027a67efb 100644 --- a/newton/_src/solvers/vbd/rigid_vbd_kernels.py +++ b/newton/_src/solvers/vbd/rigid_vbd_kernels.py @@ -12,7 +12,7 @@ - Data structures: RigidForceElementAdjacencyInfo and related structs - Device functions: Helper functions for rigid body dynamics - Utility kernels: Adjacency building -- Pre-iteration kernels: Forward integration, warmstarting, Dahl parameter computation +- Pre-iteration kernels: Forward integration, contact history restore, Dahl parameter computation - Iteration kernels: Contact accumulation, rigid body solve, dual updates - Post-iteration kernels: Velocity updates, Dahl state updates """ @@ -20,6 +20,7 @@ import warp as wp from newton._src.core.types import MAXVAL +from newton._src.geometry.hashtable import HASHTABLE_EMPTY_KEY, hashtable_find, hashtable_find_or_insert from newton._src.math import quat_velocity from newton._src.sim import JointType from newton._src.solvers.solver import integrate_rigid_body @@ -33,9 +34,6 @@ _SMALL_ANGLE_EPS = wp.constant(1.0e-7) """Small-angle threshold [rad] for guards and series expansions""" -_SMALL_ANGLE_EPS2 = _SMALL_ANGLE_EPS * _SMALL_ANGLE_EPS -"""Square of _SMALL_ANGLE_EPS""" - _DRIVE_LIMIT_MODE_NONE = wp.constant(0) _DRIVE_LIMIT_MODE_LIMIT_LOWER = wp.constant(1) _DRIVE_LIMIT_MODE_LIMIT_UPPER = wp.constant(2) @@ -50,115 +48,165 @@ _DAHL_KAPPADOT_DEADBAND = wp.constant(1.0e-6) """Deadband threshold for hysteresis direction selection""" -_NUM_CONTACT_THREADS_PER_BODY = 16 +_NUM_CONTACT_THREADS_PER_BODY = wp.constant(16) """Threads per body for contact accumulation using strided iteration""" +_CONTACT_HASH_GRID_DIM = wp.constant(128) +"""Grid dimension per axis for contact spatial hashing (128^3 = ~2M cells, fits in 21 bits)""" -# --------------------------------- -# Helper classes and device functions -# --------------------------------- +_STICK_FLAG_ANCHOR = wp.constant(1) +"""contact_stick_flag value: frozen anchor (kinematic-dynamic contacts)""" +_STICK_FLAG_DEADZONE = wp.constant(2) +"""contact_stick_flag value: anti-creep deadzone (dynamic-dynamic contacts)""" -class vec6(wp.types.vector(length=6, dtype=wp.float32)): - """Packed lower-triangular 3x3 matrix storage: [L00, L10, L11, L20, L21, L22].""" +_CONTACT_MATCH_NORMAL_DOT_THRESH = wp.constant(0.5) +"""Normal agreement gate for contact history matching (about 60 degree tolerance)""" - pass +# --------------------------------- +# Contact history hash helpers +# --------------------------------- @wp.func -def chol33(A: wp.mat33) -> vec6: - """ - Compute Cholesky factorization A = L*L^T for 3x3 SPD matrix. - - Uses packed storage for lower-triangular L to save memory and improve cache efficiency. - Packed format: [L00, L10, L11, L20, L21, L22] stores only the 6 non-zero elements. +def _pack_shape_pair_bits(s0: int, s1: int) -> wp.uint64: + """Canonical shape pair packed into upper 42 bits (21 bits each, min/max ordered).""" + lo = wp.min(s0, s1) + hi = wp.max(s0, s1) + return (wp.uint64(lo) << wp.uint64(21) | wp.uint64(hi)) << wp.uint64(21) - Algorithm: Standard column-by-column Cholesky decomposition - Column 0: L00 = sqrt(a00), L10 = a10/L00, L20 = a20/L00 - Column 1: L11 = sqrt(a11 - L10^2), L21 = (a21 - L20*L10)/L11 - Column 2: L22 = sqrt(a22 - L20^2 - L21^2) - Args: - A: Symmetric positive-definite 3x3 matrix (only lower triangle is accessed) - - Returns: - vec6: Packed lower-triangular Cholesky factor L - Layout: [L00, L10, L11, L20, L21, L22] - Represents: L = [[L00, 0, 0], - [L10, L11, 0], - [L20, L21, L22]] - - Note: Assumes A is SPD. No checking for negative square roots. - """ - # Extract lower triangle (A is symmetric, only lower half needed) - a00 = A[0, 0] - a10 = A[1, 0] - a11 = A[1, 1] - a20 = A[2, 0] - a21 = A[2, 1] - a22 = A[2, 2] - - # Column 0: Compute first column of L - L00 = wp.sqrt(a00) - L10 = a10 / L00 - L20 = a20 / L00 - - # Column 1: Compute second column of L - L11 = wp.sqrt(a11 - L10 * L10) - L21 = (a21 - L20 * L10) / L11 +@wp.func +def _compute_spatial_cell(point: wp.vec3, cell_size_inv: float) -> int: + """Quantize body-local point to 1D cell ID on a 128^3 grid.""" + gd = _CONTACT_HASH_GRID_DIM + ix = int(wp.floor(point[0] * cell_size_inv)) % gd + iy = int(wp.floor(point[1] * cell_size_inv)) % gd + iz = int(wp.floor(point[2] * cell_size_inv)) % gd + if ix < 0: + ix += gd + if iy < 0: + iy += gd + if iz < 0: + iz += gd + return ix * gd * gd + iy * gd + iz - # Column 2: Compute third column of L - L22 = wp.sqrt(a22 - L20 * L20 - L21 * L21) - # Pack into vec6: [L00, L10, L11, L20, L21, L22] - return vec6(L00, L10, L11, L20, L21, L22) +@wp.func +def _offset_cell(base_cell: int, dx: int, dy: int, dz: int) -> int: + """Compute neighbor cell ID from a base cell and 3D offset.""" + gd = _CONTACT_HASH_GRID_DIM + ix = (base_cell // (gd * gd) + dx) % gd + iy = ((base_cell // gd) % gd + dy) % gd + iz = (base_cell % gd + dz) % gd + if ix < 0: + ix += gd + if iy < 0: + iy += gd + if iz < 0: + iz += gd + return ix * gd * gd + iy * gd + iz @wp.func -def chol33_solve(Lp: vec6, b: wp.vec3) -> wp.vec3: - """ - Solve A*x = b given packed Cholesky factorization A = L*L^T. +def _build_contact_history_key(pair_bits: wp.uint64, cell: int) -> wp.uint64: + """Compose a 64-bit hash key from shape-pair bits (upper 42) and cell ID (lower 21).""" + return pair_bits | wp.uint64(cell) - Uses two-stage triangular solve: - 1. Forward substitution: L*y = b (solve for y) - 2. Backward substitution: L^T*x = y (solve for x) - This is more efficient than computing A^-1 explicitly and avoids - numerical issues from matrix inversion. +# --------------------------------- +# Helper classes and device functions +# --------------------------------- - Args: - Lp: Packed lower-triangular Cholesky factor from chol33() - Layout: [L00, L10, L11, L20, L21, L22] - b: Right-hand side vector - Returns: - vec3: Solution x to A*x = b +@wp.func +def ldlt6_solve(h_ll: wp.mat33, h_aa: wp.mat33, h_al: wp.mat33, rhs_lin: wp.vec3, rhs_ang: wp.vec3): + """Solve the 6x6 SPD block system via direct LDL^T factorization. - Complexity: 6 multiplies, 6 divides (optimal for 3x3) + Returns (x_lin, x_ang). """ - # Unpack Cholesky factor for readability - L00 = Lp[0] - L10 = Lp[1] - L11 = Lp[2] - L20 = Lp[3] - L21 = Lp[4] - L22 = Lp[5] - - # Forward substitution: L*y = b - y0 = b[0] / L00 - y1 = (b[1] - L10 * y0) / L11 - y2 = (b[2] - L20 * y0 - L21 * y1) / L22 - - # Backward substitution: L^T*x = y - x2 = y2 / L22 - x1 = (y1 - L21 * x2) / L11 - x0 = (y0 - L10 * x1 - L20 * x2) / L00 - - return wp.vec3(x0, x1, x2) + A11 = h_ll[0, 0] + A21 = h_ll[1, 0] + A22 = h_ll[1, 1] + A31 = h_ll[2, 0] + A32 = h_ll[2, 1] + A33 = h_ll[2, 2] + A41 = h_al[0, 0] + A42 = h_al[0, 1] + A43 = h_al[0, 2] + A44 = h_aa[0, 0] + A51 = h_al[1, 0] + A52 = h_al[1, 1] + A53 = h_al[1, 2] + A54 = h_aa[1, 0] + A55 = h_aa[1, 1] + A61 = h_al[2, 0] + A62 = h_al[2, 1] + A63 = h_al[2, 2] + A64 = h_aa[2, 0] + A65 = h_aa[2, 1] + A66 = h_aa[2, 2] + + # LDL^T decomposition + L21 = A21 / A11 + L31 = A31 / A11 + L41 = A41 / A11 + L51 = A51 / A11 + L61 = A61 / A11 + + D2 = A22 - L21 * L21 * A11 + + L32 = (A32 - L21 * L31 * A11) / D2 + L42 = (A42 - L21 * L41 * A11) / D2 + L52 = (A52 - L21 * L51 * A11) / D2 + L62 = (A62 - L21 * L61 * A11) / D2 + + D3 = A33 - (L31 * L31 * A11 + L32 * L32 * D2) + + L43 = (A43 - L31 * L41 * A11 - L32 * L42 * D2) / D3 + L53 = (A53 - L31 * L51 * A11 - L32 * L52 * D2) / D3 + L63 = (A63 - L31 * L61 * A11 - L32 * L62 * D2) / D3 + + D4 = A44 - (L41 * L41 * A11 + L42 * L42 * D2 + L43 * L43 * D3) + + L54 = (A54 - L41 * L51 * A11 - L42 * L52 * D2 - L43 * L53 * D3) / D4 + L64 = (A64 - L41 * L61 * A11 - L42 * L62 * D2 - L43 * L63 * D3) / D4 + + D5 = A55 - (L51 * L51 * A11 + L52 * L52 * D2 + L53 * L53 * D3 + L54 * L54 * D4) + + L65 = (A65 - L51 * L61 * A11 - L52 * L62 * D2 - L53 * L63 * D3 - L54 * L64 * D4) / D5 + + D6 = A66 - (L61 * L61 * A11 + L62 * L62 * D2 + L63 * L63 * D3 + L64 * L64 * D4 + L65 * L65 * D5) + + # Forward substitution: L y = b + y1 = rhs_lin[0] + y2 = rhs_lin[1] - L21 * y1 + y3 = rhs_lin[2] - L31 * y1 - L32 * y2 + y4 = rhs_ang[0] - L41 * y1 - L42 * y2 - L43 * y3 + y5 = rhs_ang[1] - L51 * y1 - L52 * y2 - L53 * y3 - L54 * y4 + y6 = rhs_ang[2] - L61 * y1 - L62 * y2 - L63 * y3 - L64 * y4 - L65 * y5 + + # Diagonal solve: D z = y + z1 = y1 / A11 + z2 = y2 / D2 + z3 = y3 / D3 + z4 = y4 / D4 + z5 = y5 / D5 + z6 = y6 / D6 + + # Back-substitution: L^T x = z + x6 = z6 + x5 = z5 - L65 * x6 + x4 = z4 - L54 * x5 - L64 * x6 + x3 = z3 - L43 * x4 - L53 * x5 - L63 * x6 + x2 = z2 - L32 * x3 - L42 * x4 - L52 * x5 - L62 * x6 + x1 = z1 - L21 * x2 - L31 * x3 - L41 * x4 - L51 * x5 - L61 * x6 + + return wp.vec3(x1, x2, x3), wp.vec3(x4, x5, x6) @wp.func -def cable_get_kappa(q_wp: wp.quat, q_wc: wp.quat, q_wp_rest: wp.quat, q_wc_rest: wp.quat) -> wp.vec3: +def compute_kappa(q_wp: wp.quat, q_wc: wp.quat, q_wp_rest: wp.quat, q_wc_rest: wp.quat) -> wp.vec3: """Compute cable bending curvature vector kappa in the parent frame. Kappa is the rotation vector (theta*axis) from the rest-aligned relative rotation. @@ -173,9 +221,9 @@ def cable_get_kappa(q_wp: wp.quat, q_wc: wp.quat, q_wp_rest: wp.quat, q_wc_rest: wp.vec3: Curvature vector kappa in parent frame (rotation vector form). """ # Build R_align = R_rel * R_rel_rest^T using quaternions - q_rel = wp.mul(wp.quat_inverse(q_wp), q_wc) - q_rel_rest = wp.mul(wp.quat_inverse(q_wp_rest), q_wc_rest) - q_align = wp.mul(q_rel, wp.quat_inverse(q_rel_rest)) + q_rel = wp.quat_inverse(q_wp) * q_wc + q_rel_rest = wp.quat_inverse(q_wp_rest) * q_wc_rest + q_align = q_rel * wp.quat_inverse(q_rel_rest) # Enforce shortest path (w > 0) to avoid double-cover ambiguity if q_align[3] < 0.0: @@ -209,43 +257,52 @@ def compute_right_jacobian_inverse(kappa: wp.vec3) -> wp.mat33: @wp.func -def compute_kappa_dot_analytic( - q_wp: wp.quat, - q_wc: wp.quat, - q_wp_rest: wp.quat, - q_wc_rest: wp.quat, +def compute_kappa_dot( + J_world: wp.mat33, omega_p_world: wp.vec3, omega_c_world: wp.vec3, - kappa_now: wp.vec3, ) -> wp.vec3: - """Analytical time derivative of curvature vector d(kappa)/dt in parent frame. + """Time derivative of curvature vector d(kappa)/dt in parent frame. - R_align = R_rel * R_rel_rest^T represents the rotation from rest to current configuration, - which is the same deformation measure used in cable_get_kappa. This removes the rest offset - so bending is measured relative to the undeformed state. + Exploits J_world^T = Jr_inv * R_align^T * R_wp^T, so + kappa_dot = J_world^T * (omega_c - omega_p). Args: - q_wp: Parent orientation (world). - q_wc: Child orientation (world). - q_wp_rest: Parent rest orientation (world). - q_wc_rest: Child rest orientation (world). + J_world: World-frame force Jacobian from compute_kappa_and_jacobian. omega_p_world: Parent angular velocity (world) [rad/s]. omega_c_world: Child angular velocity (world) [rad/s]. - kappa_now: Current curvature vector in parent frame. Returns: wp.vec3: Curvature rate kappa_dot in parent frame [rad/s]. """ - R_wp = wp.quat_to_matrix(q_wp) - omega_rel_parent = wp.transpose(R_wp) * (omega_c_world - omega_p_world) + return wp.transpose(J_world) * (omega_c_world - omega_p_world) + + +@wp.func +def compute_kappa_and_jacobian( + q_wp: wp.quat, + q_wc: wp.quat, + q_wp_rest: wp.quat, + q_wc_rest: wp.quat, +): + """Compute curvature vector and world-frame Jacobian from quaternion poses. + Returns: + (kappa, J_world) -- curvature vector and world-frame force Jacobian. + """ q_rel = wp.quat_inverse(q_wp) * q_wc q_rel_rest = wp.quat_inverse(q_wp_rest) * q_wc_rest - R_align = wp.quat_to_matrix(q_rel * wp.quat_inverse(q_rel_rest)) + q_align = q_rel * wp.quat_inverse(q_rel_rest) + if q_align[3] < 0.0: + q_align = wp.quat(-q_align[0], -q_align[1], -q_align[2], -q_align[3]) + axis, angle = wp.quat_to_axis_angle(q_align) + kappa = axis * angle - Jr_inv = compute_right_jacobian_inverse(kappa_now) - omega_right = wp.transpose(R_align) * omega_rel_parent - return Jr_inv * omega_right + Jr_inv = compute_right_jacobian_inverse(kappa) + R_wp = wp.quat_to_matrix(q_wp) + R_align = wp.quat_to_matrix(q_align) + J_world = R_wp * (R_align * wp.transpose(Jr_inv)) + return kappa, J_world @wp.func @@ -307,6 +364,59 @@ def build_joint_projectors( return P_lin, P_ang +@wp.func +def _average_contact_material( + ke0: float, + kd0: float, + mu0: float, + ke1: float, + kd1: float, + mu1: float, +): + """Average material properties for a contact pair. + + ke, kd: arithmetic mean. + mu: geometric mean. + """ + avg_ke = 0.5 * (ke0 + ke1) + avg_kd = 0.5 * (kd0 + kd1) + avg_mu = wp.sqrt(mu0 * mu1) + return avg_ke, avg_kd, avg_mu + + +@wp.func +def _update_dual_vec3( + C_vec: wp.vec3, + C0: wp.vec3, + alpha: float, + k: float, + lam: wp.vec3, + is_hard: int, +): + """Shared AVBD dual update for a vec3 constraint slot. + + Hard mode: stabilized constraint + lambda accumulation. + Soft mode: lambda unchanged. + + Args: + C_vec: Current constraint violation vector. + C0: Initial constraint violation snapshot for stabilization. + alpha: C0 stabilization factor. + k: Fixed penalty stiffness. + lam: Current Lagrange multiplier. + is_hard: 1 for hard (AL), 0 for soft (penalty-only). + + Returns: + wp.vec3: Updated Lagrange multiplier. + """ + if is_hard == 1: + C_stab = C_vec - alpha * C0 + lam_new = k * C_stab + lam + else: + lam_new = lam + return lam_new + + @wp.func def evaluate_angular_constraint_force_hessian( q_wp: wp.quat, @@ -316,10 +426,13 @@ def evaluate_angular_constraint_force_hessian( q_wp_prev: wp.quat, q_wc_prev: wp.quat, is_parent: bool, - k_eff: float, + penalty_k: float, P: wp.mat33, sigma0: wp.vec3, C_fric: wp.vec3, + lambda_ang: wp.vec3, + C0_ang: wp.vec3, + alpha: float, damping: float, dt: float, ): @@ -328,6 +441,9 @@ def evaluate_angular_constraint_force_hessian( Unified evaluator for all joint types. Computes constraint force and Hessian in the constrained subspace defined by the orthogonal-complement projector P. + C0 stabilization: when alpha > 0 and C0_ang is nonzero, the effective + kappa is kappa - alpha*C0_ang (initial violation snapshot). + Special cases by projector: - P = I: isotropic (CABLE bend, FIXED angular) - P = I - a*a^T: revolute (1 free angular axis) @@ -343,21 +459,13 @@ def evaluate_angular_constraint_force_hessian( """ inv_dt = 1.0 / dt - kappa_now_vec = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - kappa_perp = P * kappa_now_vec - - Jr_inv = compute_right_jacobian_inverse(kappa_now_vec) - R_wp = wp.quat_to_matrix(q_wp) - - q_rel = wp.quat_inverse(q_wp) * q_wc - q_rel_rest = wp.quat_inverse(q_wp_rest) * q_wc_rest - R_align = wp.quat_to_matrix(q_rel * wp.quat_inverse(q_rel_rest)) - - J_world = R_wp * (R_align * wp.transpose(Jr_inv)) + kappa_now_vec, J_world = compute_kappa_and_jacobian(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_stab = kappa_now_vec - alpha * C0_ang + kappa_perp = P * kappa_stab - f_local = k_eff * kappa_perp + sigma0 + f_local = penalty_k * kappa_perp + sigma0 + lambda_ang - H_local = k_eff * P + wp.mat33( + H_local = penalty_k * P + wp.mat33( C_fric[0], 0.0, 0.0, @@ -373,14 +481,12 @@ def evaluate_angular_constraint_force_hessian( omega_p_world = quat_velocity(q_wp, q_wp_prev, dt) omega_c_world = quat_velocity(q_wc, q_wc_prev, dt) - dkappa_dt_vec = compute_kappa_dot_analytic( - q_wp, q_wc, q_wp_rest, q_wc_rest, omega_p_world, omega_c_world, kappa_now_vec - ) + dkappa_dt_vec = compute_kappa_dot(J_world, omega_p_world, omega_c_world) dkappa_perp = P * dkappa_dt_vec - f_damp_local = (damping * k_eff) * dkappa_perp + f_damp_local = (damping * penalty_k) * dkappa_perp f_local = f_local + f_damp_local - k_damp = (damping * inv_dt) * k_eff + k_damp = (damping * inv_dt) * penalty_k H_local = H_local + k_damp * P H_aa = J_world * (H_local * wp.transpose(J_world)) @@ -405,6 +511,9 @@ def evaluate_linear_constraint_force_hessian( is_parent: bool, penalty_k: float, P: wp.mat33, + lambda_lin: wp.vec3, + C0_lin: wp.vec3, + alpha: float, damping: float, dt: float, ): @@ -413,6 +522,9 @@ def evaluate_linear_constraint_force_hessian( Unified evaluator for all joint types. Computes C = x_c - x_p, projects with P, and returns force/Hessian in world frame. + C0 stabilization: when alpha > 0 and C0_lin is nonzero, the effective + constraint violation is C - alpha*C0 (initial violation snapshot). + Special cases by projector: - P = I: isotropic (BALL, CABLE stretch, FIXED linear, REVOLUTE linear) - P = I - a*a^T: prismatic (1 free linear axis) @@ -436,37 +548,28 @@ def evaluate_linear_constraint_force_hessian( r = x_c - com_w C_vec = x_c - x_p - C_perp = P * C_vec - - f_attachment = penalty_k * C_perp - - rx = wp.skew(r) - K_point = penalty_k * P + C_stab = C_vec - alpha * C0_lin + C_perp = P * C_stab - H_ll = K_point - H_al = rx * K_point - H_aa = wp.transpose(rx) * K_point * rx + f_attachment = penalty_k * C_perp + lambda_lin + # Fold damping into effective stiffness: K_eff = k*(1 + d/dt)*P if damping > 0.0: + inv_dt = 1.0 / dt + K_eff = penalty_k * (1.0 + damping * inv_dt) * P + x_p_prev = wp.transform_get_translation(X_wp_prev) x_c_prev = wp.transform_get_translation(X_wc_prev) C_vec_prev = x_c_prev - x_p_prev - inv_dt = 1.0 / dt - dC_dt = (C_vec - C_vec_prev) * inv_dt - dC_dt_perp = P * dC_dt - - damping_coeff = damping * penalty_k - f_damping = damping_coeff * dC_dt_perp - f_attachment = f_attachment + f_damping - - damp_scale = damping * inv_dt - H_ll_damp = damp_scale * H_ll - H_al_damp = damp_scale * H_al - H_aa_damp = damp_scale * H_aa + dC_dt_perp = P * ((C_vec - C_vec_prev) * inv_dt) + f_attachment = f_attachment + (damping * penalty_k) * dC_dt_perp + else: + K_eff = penalty_k * P - H_ll = H_ll + H_ll_damp - H_al = H_al + H_al_damp - H_aa = H_aa + H_aa_damp + rx = wp.skew(r) + H_ll = K_eff + H_al = rx * K_eff + H_aa = wp.transpose(rx) * K_eff * rx force = f_attachment if is_parent else -f_attachment torque = wp.cross(r, force) @@ -527,47 +630,45 @@ def evaluate_rigid_contact_from_collision( body_q: wp.array[wp.transform], body_q_prev: wp.array[wp.transform], body_com: wp.array[wp.vec3], - contact_point_a_local: wp.vec3, # Local contact point on body A - contact_point_b_local: wp.vec3, # Local contact point on body B - contact_normal: wp.vec3, # Contact normal (A to B) - penetration_depth: float, # Penetration depth (> 0 when penetrating) - contact_ke: float, # Contact normal stiffness - contact_kd: float, # Contact damping coefficient - friction_mu: float, # Coulomb friction coefficient - friction_epsilon: float, # Friction regularization parameter + contact_point_a_local: wp.vec3, + contact_point_b_local: wp.vec3, + contact_normal: wp.vec3, + penetration_depth: float, + contact_ke: float, + contact_ke_t: float, + contact_kd: float, + contact_lam: wp.vec3, + friction_mu: float, + friction_epsilon: float, + hard_contact: int, dt: float, ): - """Compute contact forces and 3x3 Hessian blocks for a rigid contact pair. - - All returned forces, torques, and Hessians are in world frame. + """Compute augmented-Lagrangian contact forces and 3x3 Hessian blocks for a rigid contact pair. - Args: - body_a_index: Body A index (-1 for static/kinematic). - body_b_index: Body B index (-1 for static/kinematic). - body_q: Current body transforms (world frame). - body_q_prev: Previous body transforms (world frame) used as the - "previous" pose for finite-difference contact-relative velocity. - body_com: Body COM offsets (local body frame). - contact_point_a_local: Contact point on A (local body frame). - contact_point_b_local: Contact point on B (local body frame). - contact_normal: Unit normal from A to B (world frame). - penetration_depth: Penetration depth (> 0 when penetrating). - contact_ke: Normal stiffness. - contact_kd: Damping coefficient. - friction_mu: Coulomb friction coefficient. - friction_epsilon: Friction regularization length. - dt: Time window [s] for finite-difference contact damping/friction. + Hard contacts: ALM normal + displacement-based tangential friction with Coulomb cone clamping. + The tangential constraint is the relative tangential displacement from body_q_prev to body_q, + which correctly captures kinematic body motion. + Soft contacts: velocity-based IPC friction with scalar penalty. Returns: - tuple: (force_a, torque_a, h_ll_a, h_al_a, h_aa_a, - force_b, torque_b, h_ll_b, h_al_b, h_aa_b) - - h_ll: Linear-linear block - - h_al: Angular-linear block - - h_aa: Angular-angular block + 10-tuple: (force_a, torque_a, H_ll_a, H_al_a, H_aa_a, + force_b, torque_b, H_ll_b, H_al_b, H_aa_b) """ + lam_n = wp.dot(contact_lam, contact_normal) + + if penetration_depth <= _SMALL_LENGTH_EPS and lam_n <= 0.0: + zero_vec = wp.vec3(0.0) + zero_mat = wp.mat33(0.0) + return (zero_vec, zero_vec, zero_mat, zero_mat, zero_mat, zero_vec, zero_vec, zero_mat, zero_mat, zero_mat) + + f_n = contact_ke * penetration_depth + lam_n + if contact_ke <= 0.0: + zero_vec = wp.vec3(0.0) + zero_mat = wp.mat33(0.0) + return (zero_vec, zero_vec, zero_mat, zero_mat, zero_mat, zero_vec, zero_vec, zero_mat, zero_mat, zero_mat) + f_n = wp.max(f_n, 0.0) - # Early exit: no penetration or zero stiffness - if penetration_depth <= 0.0 or contact_ke <= 0.0: + if f_n == 0.0 and hard_contact == 0: zero_vec = wp.vec3(0.0) zero_mat = wp.mat33(0.0) return (zero_vec, zero_vec, zero_mat, zero_mat, zero_mat, zero_vec, zero_vec, zero_mat, zero_mat, zero_mat) @@ -590,88 +691,120 @@ def evaluate_rigid_contact_from_collision( X_wb_prev = body_q_prev[body_b_index] body_b_com_local = body_com[body_b_index] - # Centers of mass in world coordinates x_com_a_now = wp.transform_point(X_wa, body_a_com_local) x_com_b_now = wp.transform_point(X_wb, body_b_com_local) - # Contact points in world coordinates (current and previous) x_c_a_now = wp.transform_point(X_wa, contact_point_a_local) x_c_b_now = wp.transform_point(X_wb, contact_point_b_local) x_c_a_prev = wp.transform_point(X_wa_prev, contact_point_a_local) x_c_b_prev = wp.transform_point(X_wb_prev, contact_point_b_local) - # Contact motion for damping and friction (finite difference velocity estimation) - dx_a = x_c_a_now - x_c_a_prev # Motion of contact point on A over timestep dt - dx_b = x_c_b_now - x_c_b_prev # Motion of contact point on B over timestep dt - dx_rel = dx_b - dx_a # Relative contact motion (B relative to A) - - # Contact geometry - assume contact_normal is already unit length from collision detection - - # Normal force: f = contact_ke * penetration n_outer = wp.outer(contact_normal, contact_normal) - f_total = contact_normal * (contact_ke * penetration_depth) - K_total = contact_ke * n_outer - # Compute relative velocity for damping and friction - v_rel = dx_rel / dt + v_rel = (x_c_b_now - x_c_b_prev - x_c_a_now + x_c_a_prev) / dt v_dot_n = wp.dot(contact_normal, v_rel) - # Damping only when compressing (v_n < 0, bodies approaching) - if contact_kd > 0.0 and v_dot_n < 0.0: - damping_coeff = contact_kd * contact_ke - damping_force = -damping_coeff * v_dot_n * contact_normal - damping_hessian = (damping_coeff / dt) * n_outer - f_total = f_total + damping_force - K_total = K_total + damping_hessian - - # Normal load for friction (elastic normal only) - normal_load = contact_ke * penetration_depth - - # Friction forces (isotropic, no explicit tangent basis) - if friction_mu > 0.0 and normal_load > 0.0: - # Tangential slip (world space) - v_n = contact_normal * v_dot_n - v_t = v_rel - v_n - u = v_t * dt - eps_u = friction_epsilon * dt - - # Projected isotropic friction (no explicit tangent basis) - f_friction, K_friction = compute_projected_isotropic_friction( - friction_mu, normal_load, contact_normal, u, eps_u - ) - f_total = f_total + f_friction - K_total = K_total + K_friction + if hard_contact == 1: + f_total = contact_normal * f_n + K_total = contact_ke * n_outer + + if friction_mu > 0.0 and f_n > 0.0: + # ALM tangential friction with Coulomb cone clamping. + # Tangential constraint = relative tangential displacement (body_q_prev -> body_q). + v_t = v_rel - contact_normal * v_dot_n + tangential_disp = -(v_t * dt) + lam_t = contact_lam - contact_normal * lam_n + f_t_vec = contact_ke_t * tangential_disp + lam_t + f_t_len = wp.length(f_t_vec) + cone_limit = friction_mu * f_n + if f_t_len > cone_limit and f_t_len > 0.0: + cone_ratio = cone_limit / f_t_len + f_t_vec = f_t_vec * cone_ratio + f_total = f_total + f_t_vec + I3 = wp.identity(n=3, dtype=float) + K_total = K_total + contact_ke_t * (I3 - n_outer) + else: + # Soft contact: normal penalty + IPC velocity-based friction + f_total = contact_normal * f_n + K_total = contact_ke * n_outer + + if friction_mu > 0.0 and f_n > 0.0: + v_t = v_rel - contact_normal * v_dot_n + f_friction, K_friction = compute_projected_isotropic_friction( + friction_mu, f_n, contact_normal, v_t * dt, friction_epsilon * dt + ) + f_total = f_total + f_friction + K_total = K_total + K_friction - # Split total contact force to both bodies (Newton's 3rd law) - force_a = -f_total # Force on A (opposite to normal, pushes A away from B) - force_b = f_total # Force on B (along normal, pushes B away from A) + if contact_kd > 0.0 and v_dot_n < 0.0 and f_n > 0.0: + damping_coeff = contact_kd * contact_ke + f_total = f_total - damping_coeff * v_dot_n * contact_normal + K_total = K_total + (damping_coeff / dt) * n_outer - # Torque arms and resulting torques - r_a = x_c_a_now - x_com_a_now # Moment arm from A's COM to contact point - r_b = x_c_b_now - x_com_b_now # Moment arm from B's COM to contact point + r_a = x_c_a_now - x_com_a_now + r_b = x_c_b_now - x_com_b_now - # Angular/linear coupling using contact-point Jacobian J = [-[r]x, I] r_a_skew = wp.skew(r_a) r_a_skew_T_K = wp.transpose(r_a_skew) * K_total - torque_a = wp.cross(r_a, force_a) h_aa_a = r_a_skew_T_K * r_a_skew h_al_a = -r_a_skew_T_K - h_ll_a = K_total # Linear-linear - r_b_skew = wp.skew(r_b) r_b_skew_T_K = wp.transpose(r_b_skew) * K_total - torque_b = wp.cross(r_b, force_b) h_aa_b = r_b_skew_T_K * r_b_skew h_al_b = -r_b_skew_T_K - h_ll_b = K_total # Linear-linear + return ( + -f_total, + wp.cross(r_a, -f_total), + K_total, + h_al_a, + h_aa_a, + f_total, + wp.cross(r_b, f_total), + K_total, + h_al_b, + h_aa_b, + ) - return (force_a, torque_a, h_ll_a, h_al_a, h_aa_a, force_b, torque_b, h_ll_b, h_al_b, h_aa_b) + +@wp.func +def _compute_body_particle_contact_force( + penetration_depth: float, + n: wp.vec3, + relative_translation: wp.vec3, + ke: float, + kd: float, + mu: float, + friction_epsilon: float, + dt: float, +): + """Pure force law for body-particle contacts: normal penalty + damping + friction. + + All geometry and kinematics (penetration, normal, relative displacement) are + resolved by the caller. This function only computes the contact force and + Hessian from those scalar/vector inputs. + """ + f_n = penetration_depth * ke + force = n * f_n + hessian = ke * wp.outer(n, n) + + if wp.dot(n, relative_translation) < 0.0: + damping_coeff = kd * ke + damping_hessian = (damping_coeff / dt) * wp.outer(n, n) + hessian = hessian + damping_hessian + force = force - damping_hessian * relative_translation + + eps_u = friction_epsilon * dt + friction_force, friction_hessian = compute_projected_isotropic_friction(mu, f_n, n, relative_translation, eps_u) + force = force + friction_force + hessian = hessian + friction_hessian + + return force, hessian @wp.func -def evaluate_body_particle_contact( +def _eval_body_particle_contact( particle_index: int, particle_pos: wp.vec3, particle_prev_pos: wp.vec3, @@ -681,7 +814,6 @@ def evaluate_body_particle_contact( friction_mu: float, friction_epsilon: float, particle_radius: wp.array[float], - shape_material_mu: wp.array[float], shape_body: wp.array[int], body_q: wp.array[wp.transform], body_q_prev: wp.array[wp.transform], @@ -693,43 +825,11 @@ def evaluate_body_particle_contact( contact_normal: wp.array[wp.vec3], dt: float, ): - """ - Evaluate particle-rigid body contact force and Hessian (on particle side). - - Computes contact forces and Hessians for a particle interacting with a rigid body shape. - The function is agnostic to whether the rigid body is static, kinematic, or dynamic. - - Contact model: - - Normal: Linear spring-damper (stiffness: body_particle_contact_ke, damping: body_particle_contact_kd) - - Friction: 3D projector-based Coulomb friction with IPC regularization - - Normal direction: Points from rigid surface towards particle (into particle) + """Particle-rigid contact force/Hessian - resolves geometry from arrays then + delegates to ``_compute_body_particle_contact_force``. - Args: - particle_index: Index of the particle - particle_pos: Current particle position (world frame) - particle_prev_pos: Previous particle position (world frame) used as the - "previous" position for finite-difference contact-relative velocity. - contact_index: Index in the body-particle contact arrays - body_particle_contact_ke: Contact stiffness (model-level or AVBD adaptive) - body_particle_contact_kd: Contact damping (model-level or AVBD averaged) - friction_mu: Friction coefficient (model-level or AVBD averaged) - friction_epsilon: Friction regularization distance - particle_radius: Array of particle radii - shape_material_mu: Array of shape friction coefficients - shape_body: Array mapping shape index to body index - body_q: Current body transforms - body_q_prev: Previous body transforms (for finite-difference body - velocity when available) - body_qd: Body spatial velocities (fallback when no previous pose is provided) - body_com: Body centers of mass (local frame) - contact_shape: Array of shape indices for each soft contact - contact_body_pos: Array of contact points (local to shape) - contact_body_vel: Array of contact velocities (local frame) - contact_normal: Array of contact normals (world frame, from rigid to particle) - dt: Time window [s] used for finite-difference damping/friction. - - Returns: - tuple[wp.vec3, wp.mat33]: (force, Hessian) on the particle (world frame) + Prefer calling ``_compute_body_particle_contact_force`` directly when the + caller already has the contact geometry and relative displacement. """ shape_index = contact_shape[contact_index] body_index = shape_body[shape_index] @@ -740,66 +840,96 @@ def evaluate_body_particle_contact( 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[contact_index]) - n = contact_normal[contact_index] penetration_depth = -(wp.dot(n, particle_pos - bx) - particle_radius[particle_index]) if penetration_depth > 0.0: - body_contact_force_norm = penetration_depth * body_particle_contact_ke - body_contact_force = n * body_contact_force_norm - body_contact_hessian = body_particle_contact_ke * wp.outer(n, n) - - # Combine body-particle friction and shape material friction using geometric mean. - mu = wp.sqrt(friction_mu * shape_material_mu[shape_index]) - dx = particle_pos - particle_prev_pos - if wp.dot(n, dx) < 0.0: - # Damping coefficient is scaled by contact stiffness (consistent with rigid-rigid) - damping_coeff = body_particle_contact_kd * body_particle_contact_ke - damping_hessian = (damping_coeff / dt) * wp.outer(n, n) - body_contact_hessian = body_contact_hessian + damping_hessian - body_contact_force = body_contact_force - damping_hessian * dx - - # body velocity if body_q_prev: - # if body_q_prev is available, compute velocity using finite difference method - # this is more accurate for simulating static friction X_wb_prev = wp.transform_identity() if body_index >= 0: X_wb_prev = body_q_prev[body_index] bx_prev = wp.transform_point(X_wb_prev, contact_body_pos[contact_index]) bv = (bx - bx_prev) / dt + wp.transform_vector(X_wb, contact_body_vel[contact_index]) - else: - # otherwise use the instantaneous velocity r = bx - wp.transform_point(X_wb, X_com) body_v_s = wp.spatial_vector() if body_index >= 0: body_v_s = body_qd[body_index] - body_w = wp.spatial_bottom(body_v_s) body_v = wp.spatial_top(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[contact_index]) relative_translation = dx - bv * dt - # Friction using 3D projector approach (consistent with rigid-rigid contacts) - eps_u = friction_epsilon * dt - friction_force, friction_hessian = compute_projected_isotropic_friction( - mu, body_contact_force_norm, n, relative_translation, eps_u + return _compute_body_particle_contact_force( + penetration_depth, + n, + relative_translation, + body_particle_contact_ke, + body_particle_contact_kd, + friction_mu, + friction_epsilon, + dt, ) - body_contact_force = body_contact_force + friction_force - body_contact_hessian = body_contact_hessian + friction_hessian else: - body_contact_force = wp.vec3(0.0, 0.0, 0.0) - body_contact_hessian = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + return wp.vec3(0.0), wp.mat33(0.0) - return body_contact_force, body_contact_hessian + +@wp.func +def evaluate_body_particle_contact( + particle_index: int, + particle_pos: wp.vec3, + particle_prev_pos: wp.vec3, + contact_index: int, + body_particle_contact_ke: float, + body_particle_contact_kd: float, + friction_mu: float, + friction_epsilon: float, + particle_radius: wp.array[float], + shape_material_mu: wp.array[float], + shape_body: wp.array[int], + body_q: wp.array[wp.transform], + body_q_prev: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], + body_com: wp.array[wp.vec3], + contact_shape: wp.array[int], + contact_body_pos: wp.array[wp.vec3], + contact_body_vel: wp.array[wp.vec3], + contact_normal: wp.array[wp.vec3], + dt: float, +): + """Particle-rigid contact force/Hessian with per-shape mu mixing. + + VBD rigid-side uses ``_eval_body_particle_contact`` directly (mu is + pre-averaged per contact). This wrapper is kept for other solvers + that pass raw mu and need per-shape mixing. + """ + shape_index = contact_shape[contact_index] + mixed_mu = wp.sqrt(friction_mu * shape_material_mu[shape_index]) + return _eval_body_particle_contact( + particle_index, + particle_pos, + particle_prev_pos, + contact_index, + body_particle_contact_ke, + body_particle_contact_kd, + mixed_mu, + friction_epsilon, + particle_radius, + shape_body, + body_q, + body_q_prev, + body_qd, + body_com, + contact_shape, + contact_body_pos, + contact_body_vel, + contact_normal, + dt, + ) @wp.func @@ -856,6 +986,14 @@ def resolve_drive_limit_mode( has_drive: bool, has_limits: bool, ): + """Resolve drive/limit priority and compute position error [m or rad]. + + Limits take precedence: if q is outside [lower, upper], the active limit + wins. Otherwise the drive engages with target clamped to the limit range. + + Returns: + (mode, err_pos) -- active mode constant and signed position error. + """ mode = _DRIVE_LIMIT_MODE_NONE err_pos = float(0.0) drive_target = target_pos @@ -873,23 +1011,6 @@ def resolve_drive_limit_mode( return mode, err_pos -@wp.func -def compute_kappa_and_jacobian( - q_wp: wp.quat, - q_wc: wp.quat, - q_wp_rest: wp.quat, - q_wc_rest: wp.quat, -): - kappa = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - Jr_inv = compute_right_jacobian_inverse(kappa) - R_wp = wp.quat_to_matrix(q_wp) - q_rel = wp.quat_inverse(q_wp) * q_wc - q_rel_rest = wp.quat_inverse(q_wp_rest) * q_wc_rest - R_align = wp.quat_to_matrix(q_rel * wp.quat_inverse(q_rel_rest)) - J_world = R_wp * (R_align * wp.transpose(Jr_inv)) - return kappa, J_world - - @wp.func def apply_angular_drive_limit_torque( a: wp.vec3, @@ -898,10 +1019,14 @@ def apply_angular_drive_limit_torque( f_scalar: float, H_scalar: float, ): - f_local = f_scalar * a - H_local = H_scalar * wp.outer(a, a) - tau = J_world * f_local - Haa = J_world * (H_local * wp.transpose(J_world)) + """Rank-1 angular drive/limit torque and Hessian along local axis a. + + Maps scalar spring-damper (f_scalar, H_scalar) through J_world to + world-frame torque and H_aa. + """ + Ja = J_world * a + tau = f_scalar * Ja + Haa = H_scalar * wp.outer(Ja, Ja) if not is_parent: tau = -tau return tau, Haa @@ -915,18 +1040,27 @@ def apply_linear_drive_limit_force( f_scalar: float, H_scalar: float, ): + """Rank-1 linear drive/limit force and Hessian along world axis. + + Maps scalar spring-damper (f_scalar, H_scalar) to world-frame force, + torque, and Hessian blocks (H_ll, H_al, H_aa) via the moment arm r. + """ f_attachment = f_scalar * axis_w - aa = wp.outer(axis_w, axis_w) - K_point = H_scalar * aa - rx = wp.skew(r) - Hll = K_point - Hal = rx * K_point - Haa = wp.transpose(rx) * K_point * rx + ra = wp.cross(r, axis_w) + Hll = H_scalar * wp.outer(axis_w, axis_w) + Hal = H_scalar * wp.outer(ra, axis_w) + Haa = H_scalar * wp.outer(ra, ra) force = f_attachment if is_parent else -f_attachment torque = wp.cross(r, force) return force, torque, Hll, Hal, Haa +@wp.func +def _zero_force_hessian(): + """Zero (force, torque, H_ll, H_al, H_aa) tuple for early-exit paths.""" + return wp.vec3(0.0), wp.vec3(0.0), wp.mat33(0.0), wp.mat33(0.0), wp.mat33(0.0) + + @wp.func def evaluate_joint_force_hessian( body_index: int, @@ -958,6 +1092,12 @@ def evaluate_joint_force_hessian( joint_limit_upper: wp.array[float], joint_limit_ke: wp.array[float], joint_limit_kd: wp.array[float], + joint_lambda_lin: wp.array[wp.vec3], + joint_lambda_ang: wp.array[wp.vec3], + joint_C0_lin: wp.array[wp.vec3], + joint_C0_ang: wp.array[wp.vec3], + joint_is_hard: wp.array[wp.int32], + avbd_alpha: float, joint_dof_dim: wp.array2d[int], joint_rest_angle: wp.array[float], dt: float, @@ -968,15 +1108,16 @@ def evaluate_joint_force_hessian( Uses unified projector-based constraint evaluators for all joint types. Indexing: - ``joint_constraint_start[j]`` is a solver-owned start offset into the per-constraint arrays - (``joint_penalty_k``, ``joint_penalty_kd``). The layout is: - - ``JointType.CABLE``: 2 scalars -> [stretch (linear), bend (angular)] - - ``JointType.BALL``: 1 scalar -> [linear] - - ``JointType.FIXED``: 2 scalars -> [linear, angular] - - ``JointType.REVOLUTE``: 3 scalars -> [linear, angular, ang_drive_limit] - - ``JointType.PRISMATIC``: 3 scalars -> [linear, angular, lin_drive_limit] - - ``JointType.D6``: 2 + lin_count + ang_count scalars - -> [linear, angular, lin_dl_0, ..., ang_dl_0, ...] + joint_constraint_start[j] is a solver-owned start offset into the per-constraint + arrays (joint_penalty_k, joint_penalty_kd). Layout per joint type: + - CABLE: 2 scalars -> [stretch, bend] + - BALL: 1 scalar -> [linear] + - FIXED: 2 scalars -> [linear, angular] + - REVOLUTE: 3 scalars -> [linear, angular, ang_drive_limit] + - PRISMATIC: 3 scalars -> [linear, angular, lin_drive_limit] + - D6: 2 + lin_count + ang_count scalars -> [linear, angular, per-DOF drive/limit] + Drive/limit slots use AVBD-ramped stiffness via min(avbd_ke, model_ke). + Drive/limit forces remain penalty-only (no lambda or C0 state). """ jt = joint_type[joint_index] if ( @@ -987,15 +1128,15 @@ def evaluate_joint_force_hessian( and jt != JointType.PRISMATIC and jt != JointType.D6 ): - return wp.vec3(0.0), wp.vec3(0.0), wp.mat33(0.0), wp.mat33(0.0), wp.mat33(0.0) + return _zero_force_hessian() if not joint_enabled[joint_index]: - return wp.vec3(0.0), wp.vec3(0.0), wp.mat33(0.0), wp.mat33(0.0), wp.mat33(0.0) + return _zero_force_hessian() parent_index = joint_parent[joint_index] child_index = joint_child[joint_index] if body_index != child_index and (parent_index < 0 or body_index != parent_index): - return wp.vec3(0.0), wp.vec3(0.0), wp.mat33(0.0), wp.mat33(0.0), wp.mat33(0.0) + return _zero_force_hessian() is_parent_body = parent_index >= 0 and body_index == parent_index @@ -1036,7 +1177,28 @@ def evaluate_joint_force_hessian( q_wc_prev = wp.transform_get_rotation(X_wc_prev) P_I = wp.identity(3, float) - no_dahl = wp.vec3(0.0) + + # Hard/soft AL gating for the linear structural slot (slot 0) + lin_lambda = wp.vec3(0.0) + lin_C0 = wp.vec3(0.0) + lin_alpha = float(0.0) + if joint_is_hard[c_start] == 1: + lin_lambda = joint_lambda_lin[joint_index] + lin_C0 = joint_C0_lin[joint_index] + lin_alpha = avbd_alpha + + # Hard/soft AL gating for the angular structural slot (slot 1) + ang_lambda = wp.vec3(0.0) + ang_C0 = wp.vec3(0.0) + ang_alpha = float(0.0) + ang_hard = 0 + if jt != JointType.BALL: + ang_hard = joint_is_hard[c_start + 1] + + if ang_hard == 1: + ang_lambda = joint_lambda_ang[joint_index] + ang_C0 = joint_C0_ang[joint_index] + ang_alpha = avbd_alpha if jt == JointType.CABLE: k_stretch = joint_penalty_k[c_start] @@ -1051,8 +1213,12 @@ def evaluate_joint_force_hessian( total_H_aa = wp.mat33(0.0) if k_bend > 0.0: - sigma0 = joint_sigma_start[joint_index] - C_fric = joint_C_fric[joint_index] + if ang_hard == 1: + sigma0 = wp.vec3(0.0) + C_fric = wp.vec3(0.0) + else: + sigma0 = joint_sigma_start[joint_index] + C_fric = joint_C_fric[joint_index] bend_torque, bend_H_aa, _bend_kappa, _bend_J = evaluate_angular_constraint_force_hessian( q_wp, q_wc, @@ -1065,6 +1231,9 @@ def evaluate_joint_force_hessian( P_I, sigma0, C_fric, + ang_lambda, + ang_C0, + ang_alpha, kd_bend, dt, ) @@ -1084,6 +1253,9 @@ def evaluate_joint_force_hessian( is_parent_body, k_stretch, P_I, + lin_lambda, + lin_C0, + lin_alpha, kd_stretch, dt, ) @@ -1111,10 +1283,13 @@ def evaluate_joint_force_hessian( is_parent_body, k, P_I, + lin_lambda, + lin_C0, + lin_alpha, damping, dt, ) - return wp.vec3(0.0), wp.vec3(0.0), wp.mat33(0.0), wp.mat33(0.0), wp.mat33(0.0) + return _zero_force_hessian() elif jt == JointType.FIXED: k_lin = joint_penalty_k[c_start + 0] @@ -1132,6 +1307,9 @@ def evaluate_joint_force_hessian( is_parent_body, k_lin, P_I, + lin_lambda, + lin_C0, + lin_alpha, kd_lin, dt, ) @@ -1155,8 +1333,11 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_I, - no_dahl, - no_dahl, + wp.vec3(0.0), + wp.vec3(0.0), + ang_lambda, + ang_C0, + ang_alpha, kd_ang, dt, ) @@ -1186,6 +1367,9 @@ def evaluate_joint_force_hessian( is_parent_body, k_lin, P_lin, + lin_lambda, + lin_C0, + lin_alpha, kd_lin, dt, ) @@ -1214,8 +1398,11 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_ang, - no_dahl, - no_dahl, + wp.vec3(0.0), + wp.vec3(0.0), + ang_lambda, + ang_C0, + ang_alpha, kd_ang, dt, ) @@ -1224,7 +1411,7 @@ def evaluate_joint_force_hessian( t_ang = wp.vec3(0.0) Haa_ang = wp.mat33(0.0) - # Drive + limits on free angular DOF + # Drive + limits on free angular DOF (AVBD slot c_start + 2) dof_idx = qd_start model_drive_ke = joint_target_ke[dof_idx] drive_kd = joint_target_kd[dof_idx] @@ -1255,7 +1442,7 @@ def evaluate_joint_force_hessian( theta_abs = theta + joint_rest_angle[dof_idx] omega_p = quat_velocity(q_wp, q_wp_prev, dt) omega_c = quat_velocity(q_wc, q_wc_prev, dt) - dkappa_dt = compute_kappa_dot_analytic(q_wp, q_wc, q_wp_rest, q_wc_rest, omega_p, omega_c, kappa) + dkappa_dt = compute_kappa_dot(J_world, omega_p, omega_c) dtheta_dt = wp.dot(dkappa_dt, a) mode, err_pos = resolve_drive_limit_mode(theta_abs, target_pos, lim_lower, lim_upper, has_drive, has_limits) @@ -1298,6 +1485,9 @@ def evaluate_joint_force_hessian( is_parent_body, k_lin, P_lin, + lin_lambda, + lin_C0, + lin_alpha, kd_lin, dt, ) @@ -1321,8 +1511,11 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_ang, - no_dahl, - no_dahl, + wp.vec3(0.0), + wp.vec3(0.0), + ang_lambda, + ang_C0, + ang_alpha, kd_ang, dt, ) @@ -1330,7 +1523,7 @@ def evaluate_joint_force_hessian( t_ang = wp.vec3(0.0) Haa_ang = wp.mat33(0.0) - # Drive + limits on free linear DOF + # Drive + limits on free linear DOF (AVBD slot c_start + 2) dof_idx = qd_start model_drive_ke = joint_target_ke[dof_idx] drive_kd = joint_target_kd[dof_idx] @@ -1433,6 +1626,9 @@ def evaluate_joint_force_hessian( is_parent_body, k_lin, P_lin, + lin_lambda, + lin_C0, + lin_alpha, kd_lin, dt, ) @@ -1461,8 +1657,11 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_ang, - no_dahl, - no_dahl, + wp.vec3(0.0), + wp.vec3(0.0), + ang_lambda, + ang_C0, + ang_alpha, kd_ang, dt, ) @@ -1552,7 +1751,7 @@ def evaluate_joint_force_hessian( omega_p = quat_velocity(q_wp, q_wp_prev, dt) omega_c = quat_velocity(q_wc, q_wc_prev, dt) - dkappa_dt = compute_kappa_dot_analytic(q_wp, q_wc, q_wp_rest, q_wc_rest, omega_p, omega_c, kappa) + dkappa_dt = compute_kappa_dot(J_world, omega_p, omega_c) for ai in range(3): if ai < ang_count: @@ -1603,7 +1802,7 @@ def evaluate_joint_force_hessian( return total_force, total_torque, total_H_ll, total_H_al, total_H_aa - return wp.vec3(0.0), wp.vec3(0.0), wp.mat33(0.0), wp.mat33(0.0), wp.mat33(0.0) + return _zero_force_hessian() # ----------------------------- @@ -1672,14 +1871,10 @@ def forward_step_rigid_bodies( body_q: wp.array[wp.transform], body_qd: wp.array[wp.spatial_vector], body_inertia_q: wp.array[wp.transform], - body_q_prev: wp.array[wp.transform], ): """ Forward integration step for rigid bodies in the AVBD/VBD solver. - Snapshots ``body_q_prev`` for dynamic bodies only. Kinematic bodies keep - the previous step's pose so contact friction sees correct velocity. - Args: dt: Time step [s]. gravity: Gravity vector array (world frame). @@ -1692,22 +1887,17 @@ def forward_step_rigid_bodies( body_q: Body transforms (input: start-of-step pose, output: integrated pose). body_qd: Body velocities (input: start-of-step velocity, output: integrated velocity). body_inertia_q: Inertial target body transforms for the AVBD solve (output). - body_q_prev: Previous body transforms (output, dynamic bodies only). """ tid = wp.tid() q_current = body_q[tid] # Early exit for kinematic bodies (inv_mass == 0). - # Do not snapshot body_q_prev here: kinematic bodies need body_q_prev from previous step. inv_m = body_inv_mass[tid] if inv_m == 0.0: body_inertia_q[tid] = q_current return - # Snapshot current pose as previous before integration (dynamic bodies only). - body_q_prev[tid] = q_current - # Read body state (only for dynamic bodies) qd_current = body_qd[tid] f_current = body_f[tid] @@ -1746,22 +1936,11 @@ def build_body_body_contact_lists( body_contact_buffer_pre_alloc: int, body_contact_counts: wp.array[wp.int32], body_contact_indices: wp.array[wp.int32], + body_contact_overflow_max: wp.array[wp.int32], ): """ Build per-body contact lists for body-centric per-color contact evaluation. - - Each contact index t_id is appended to both bodies' lists (shape0's body and - shape1's body), enabling efficient lookup of all contacts involving a given body - during per-body solve iterations. - - Notes: - - body_contact_counts[b] is reset to 0 on the host before launch and - atomically incremented here; consumers must only read the first - body_contact_counts[b] indices for each body. - - If a body has more than body_contact_buffer_pre_alloc contacts, extra indices - are dropped (overflow is safely ignored here). - - body_contact_indices is not cleared each step; only the prefix defined - by body_contact_counts is considered valid. + Tracks overflow into body_contact_overflow_max for diagnostics. """ t_id = wp.tid() if t_id >= rigid_contact_count[0]: @@ -1772,17 +1951,19 @@ def build_body_body_contact_lists( b0 = shape_body[s0] if s0 >= 0 else -1 b1 = shape_body[s1] if s1 >= 0 else -1 - # Add contact to body0's list if b0 >= 0: idx = wp.atomic_add(body_contact_counts, b0, 1) if idx < body_contact_buffer_pre_alloc: body_contact_indices[b0 * body_contact_buffer_pre_alloc + idx] = t_id + else: + wp.atomic_max(body_contact_overflow_max, 0, idx + 1) - # Add contact to body1's list if b1 >= 0: idx = wp.atomic_add(body_contact_counts, b1, 1) if idx < body_contact_buffer_pre_alloc: body_contact_indices[b1 * body_contact_buffer_pre_alloc + idx] = t_id + else: + wp.atomic_max(body_contact_overflow_max, 0, idx + 1) @wp.kernel @@ -1793,22 +1974,11 @@ def build_body_particle_contact_lists( body_particle_contact_buffer_pre_alloc: int, body_particle_contact_counts: wp.array[wp.int32], body_particle_contact_indices: wp.array[wp.int32], + body_particle_contact_overflow_max: wp.array[wp.int32], ): """ - Build per-body contact lists for body-particle (particle-rigid) contacts. - - Each body-particle contact index tid (from the external contacts.soft_contact_* - arrays) is appended to the list of the rigid body associated with - body_particle_contact_shape[tid] via shape_body. This enables efficient - per-body processing of particle-rigid contacts during the rigid-body solve - (e.g., Gauss-Seidel color sweeps). - - Notes: - - body_particle_contact_counts[b] must be zeroed on the host before launch. - - Only the first body_particle_contact_counts[b] entries in the flat - body_particle_contact_indices array are considered valid for body b. - - If a body has more than body_particle_contact_buffer_pre_alloc body-particle - contacts, extra indices are dropped (overflow is ignored safely). + Build per-body contact lists for body-particle contacts. + Tracks overflow into body_particle_contact_overflow_max for diagnostics. """ tid = wp.tid() if tid >= body_particle_contact_count[0]: @@ -1823,77 +1993,449 @@ def build_body_particle_contact_lists( idx = wp.atomic_add(body_particle_contact_counts, body, 1) if idx < body_particle_contact_buffer_pre_alloc: body_particle_contact_indices[body * body_particle_contact_buffer_pre_alloc + idx] = tid + else: + wp.atomic_max(body_particle_contact_overflow_max, 0, idx + 1) @wp.kernel -def warmstart_joints( - # Inputs - joint_penalty_k_max: wp.array[float], - joint_penalty_k_min: wp.array[float], +def check_contact_overflow( + overflow_max: wp.array[wp.int32], + buffer_size: int, + contact_type: int, +): + """Print a warning if per-body contact buffer overflowed. Launched with dim=1.""" + omax = overflow_max[0] + if omax > buffer_size: + if contact_type == 0: + wp.printf( + "Warning: Per-body rigid contact buffer overflowed %d > %d.\n", + omax, + buffer_size, + ) + else: + wp.printf( + "Warning: Per-body particle contact buffer overflowed %d > %d.\n", + omax, + buffer_size, + ) + + +@wp.kernel +def step_joint_C0_lambda( + joint_enabled: wp.array[bool], + joint_parent: wp.array[int], + joint_child: wp.array[int], + joint_X_p: wp.array[wp.transform], + joint_X_c: wp.array[wp.transform], + body_q_prev: wp.array[wp.transform], + body_q_rest: wp.array[wp.transform], + joint_constraint_start: wp.array[wp.int32], + joint_constraint_dim: wp.array[wp.int32], + joint_is_hard: wp.array[wp.int32], gamma: float, - # Input/output joint_penalty_k: wp.array[float], + joint_penalty_k_min: wp.array[float], + joint_penalty_k_max: wp.array[float], + joint_C0_lin: wp.array[wp.vec3], + joint_C0_ang: wp.array[wp.vec3], + joint_lambda_lin: wp.array[wp.vec3], + joint_lambda_ang: wp.array[wp.vec3], ): - """ - Warm-start per-scalar-constraint penalty stiffness for rigid joint constraints (runs once per step). + """Per-step joint AVBD maintenance: k decay + C0 snapshot + lambda decay. - Algorithm (per constraint scalar i): - - Decay previous penalty: k <- gamma * k (typically gamma in [0, 1]) - - Clamp to [joint_penalty_k_min[i], joint_penalty_k_max[i]] so k_min <= k <= k_max always. + Sole owner of all joint decay. Runs every step. """ - i = wp.tid() + j = wp.tid() + c_start = int(joint_constraint_start[j]) + c_dim = int(joint_constraint_dim[j]) + + # K decay runs unconditionally (even for disabled joints). + for s in range(c_dim): + idx = c_start + s + joint_penalty_k[idx] = wp.clamp( + gamma * joint_penalty_k[idx], joint_penalty_k_min[idx], joint_penalty_k_max[idx] + ) + + child = joint_child[j] + if not joint_enabled[j] or c_dim == 0 or child < 0: + joint_C0_lin[j] = wp.vec3(0.0) + joint_C0_ang[j] = wp.vec3(0.0) + joint_lambda_lin[j] = wp.vec3(0.0) + joint_lambda_ang[j] = wp.vec3(0.0) + return + + lin_hard = joint_is_hard[c_start] + ang_hard = 0 + if c_dim > 1: + ang_hard = joint_is_hard[c_start + 1] + + if lin_hard == 1 or ang_hard == 1: + parent = joint_parent[j] + if parent >= 0: + X_wp = body_q_prev[parent] * joint_X_p[j] + else: + X_wp = joint_X_p[j] + X_wc = body_q_prev[child] * joint_X_c[j] + + if lin_hard == 1: + x_p = wp.transform_get_translation(X_wp) + x_c = wp.transform_get_translation(X_wc) + joint_C0_lin[j] = x_c - x_p + joint_lambda_lin[j] = joint_lambda_lin[j] * gamma + else: + joint_C0_lin[j] = wp.vec3(0.0) + joint_lambda_lin[j] = wp.vec3(0.0) - k_prev = joint_penalty_k[i] - k_new = wp.clamp(gamma * k_prev, joint_penalty_k_min[i], joint_penalty_k_max[i]) - joint_penalty_k[i] = k_new + if ang_hard == 1: + if parent >= 0: + X_wp_rest = body_q_rest[parent] * joint_X_p[j] + else: + X_wp_rest = joint_X_p[j] + X_wc_rest = body_q_rest[child] * joint_X_c[j] + q_wp = wp.transform_get_rotation(X_wp) + q_wc = wp.transform_get_rotation(X_wc) + q_wp_rest = wp.transform_get_rotation(X_wp_rest) + q_wc_rest = wp.transform_get_rotation(X_wc_rest) + joint_C0_ang[j] = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + joint_lambda_ang[j] = joint_lambda_ang[j] * gamma + else: + joint_C0_ang[j] = wp.vec3(0.0) + joint_lambda_ang[j] = wp.vec3(0.0) + else: + joint_C0_lin[j] = wp.vec3(0.0) + joint_C0_ang[j] = wp.vec3(0.0) + joint_lambda_lin[j] = wp.vec3(0.0) + joint_lambda_ang[j] = wp.vec3(0.0) @wp.kernel -def warmstart_body_body_contacts( +def init_body_body_contact_materials( rigid_contact_count: wp.array[int], rigid_contact_shape0: wp.array[int], rigid_contact_shape1: wp.array[int], shape_material_ke: wp.array[float], shape_material_kd: wp.array[float], shape_material_mu: wp.array[float], - k_start_body_contact: float, + k_start: float, # Outputs contact_penalty_k: wp.array[float], - contact_material_ke: wp.array[float], contact_material_kd: wp.array[float], contact_material_mu: wp.array[float], + contact_material_ke: wp.array[float], ): - """ - Warm-start contact penalties and cache material properties. + """Cold-start body-body contact penalties and cache material properties. - Computes averaged material properties for each rigid contact and resets the - AVBD penalty to a bounded initial value based on `k_start_body_contact` and the - material stiffness. + Averages both shapes' material. Penalty is seeded at ``min(k_start, avg_ke)`` + when ramping (k_start >= 0) or at ``avg_ke`` when fixed-k (k_start < 0). """ i = wp.tid() if i >= rigid_contact_count[0]: return - # Read shape indices shape_id_0 = rigid_contact_shape0[i] shape_id_1 = rigid_contact_shape1[i] - # Cache averaged material properties (arithmetic mean for stiffness/damping, geometric for friction) - avg_ke = 0.5 * (shape_material_ke[shape_id_0] + shape_material_ke[shape_id_1]) - avg_kd = 0.5 * (shape_material_kd[shape_id_0] + shape_material_kd[shape_id_1]) - avg_mu = wp.sqrt(shape_material_mu[shape_id_0] * shape_material_mu[shape_id_1]) + avg_ke, avg_kd, avg_mu = _average_contact_material( + shape_material_ke[shape_id_0], + shape_material_kd[shape_id_0], + shape_material_mu[shape_id_0], + shape_material_ke[shape_id_1], + shape_material_kd[shape_id_1], + shape_material_mu[shape_id_1], + ) + contact_material_kd[i] = avg_kd + contact_material_mu[i] = avg_mu contact_material_ke[i] = avg_ke + + k_floor = avg_ke if k_start < 0.0 else wp.min(k_start, avg_ke) + contact_penalty_k[i] = k_floor + + +@wp.kernel +def init_body_body_contacts_avbd( + # Dimensioning + rigid_contact_count: wp.array[int], + # Constraint data + rigid_contact_shape0: wp.array[int], + rigid_contact_shape1: wp.array[int], + rigid_contact_point0: wp.array[wp.vec3], + rigid_contact_point1: wp.array[wp.vec3], + rigid_contact_normal: wp.array[wp.vec3], + # Material + shape_material_ke: wp.array[float], + shape_material_kd: wp.array[float], + shape_material_mu: wp.array[float], + hard_contacts: int, + # Cross-step state + ht_keys: wp.array[wp.uint64], + history_point0: wp.array[wp.vec3], + history_point1: wp.array[wp.vec3], + history_lambda: wp.array[wp.vec3], + history_normal: wp.array[wp.vec3], + history_stick_flag: wp.array[wp.int32], + history_penalty_k: wp.array[float], + # Scalar parameters + cell_size_inv: float, + match_tolerance_sq: float, + k_start: float, + # Outputs + contact_penalty_k: wp.array[float], + contact_lambda: wp.array[wp.vec3], + contact_material_kd: wp.array[float], + contact_material_mu: wp.array[float], + contact_material_ke: wp.array[float], +): + """Restore body-body contact state from hash table history. + + For hard contacts: restores lambda (rotated from old to new contact frame), + penalty_k, and stick-anchor points. + For soft contacts: only penalty_k restoration is meaningful (lambda and + stick_flag are stored but inert in the soft penalty path). + C0 and decay are handled by step_body_body_contact_C0_lambda. + """ + i = wp.tid() + if i >= rigid_contact_count[0]: + return + + s0 = rigid_contact_shape0[i] + s1 = rigid_contact_shape1[i] + + avg_ke, avg_kd, avg_mu = _average_contact_material( + shape_material_ke[s0], + shape_material_kd[s0], + shape_material_mu[s0], + shape_material_ke[s1], + shape_material_kd[s1], + shape_material_mu[s1], + ) contact_material_kd[i] = avg_kd contact_material_mu[i] = avg_mu - # Reset contact penalty to k_start every frame because contact indices are not persistent across frames. - k_new = wp.min(k_start_body_contact, avg_ke) - contact_penalty_k[i] = k_new + p0 = rigid_contact_point0[i] + p1 = rigid_contact_point1[i] + n_new = rigid_contact_normal[i] + pair_bits = _pack_shape_pair_bits(s0, s1) + base_cell = _compute_spatial_cell(p0, cell_size_inv) + + best_dist_sq = match_tolerance_sq + best_slot = int(-1) + + center_key = _build_contact_history_key(pair_bits, base_cell) + center_slot = hashtable_find(center_key, ht_keys) + + if center_slot >= 0: + dp0 = p0 - history_point0[center_slot] + dp1 = p1 - history_point1[center_slot] + dist_sq = wp.dot(dp0, dp0) + wp.dot(dp1, dp1) + n_dot = wp.dot(n_new, history_normal[center_slot]) + if dist_sq < best_dist_sq and n_dot > _CONTACT_MATCH_NORMAL_DOT_THRESH: + best_dist_sq = dist_sq + best_slot = center_slot + + if best_slot < 0: + for dx in range(-1, 2): + for dy in range(-1, 2): + for dz in range(-1, 2): + if dx == 0 and dy == 0 and dz == 0: + continue + cell = _offset_cell(base_cell, dx, dy, dz) + key = _build_contact_history_key(pair_bits, cell) + slot = hashtable_find(key, ht_keys) + if slot >= 0: + dp0 = p0 - history_point0[slot] + dp1 = p1 - history_point1[slot] + dist_sq = wp.dot(dp0, dp0) + wp.dot(dp1, dp1) + n_dot = wp.dot(n_new, history_normal[slot]) + if dist_sq < best_dist_sq and n_dot > _CONTACT_MATCH_NORMAL_DOT_THRESH: + best_dist_sq = dist_sq + best_slot = slot + + contact_material_ke[i] = avg_ke + k_floor = avg_ke if k_start < 0.0 else wp.min(k_start, avg_ke) + + if best_slot >= 0: + contact_penalty_k[i] = wp.clamp(history_penalty_k[best_slot], k_floor, avg_ke) + lam_hist = history_lambda[best_slot] + if hard_contacts == 1: + n_old = history_normal[best_slot] + lam_n = wp.dot(lam_hist, n_old) + lam_t_old = lam_hist - n_old * lam_n + lam_t_new = lam_t_old - n_new * wp.dot(lam_t_old, n_new) + contact_lambda[i] = n_new * lam_n + lam_t_new + else: + contact_lambda[i] = lam_hist + + if hard_contacts == 1 and history_stick_flag[best_slot] == _STICK_FLAG_ANCHOR: + rigid_contact_point0[i] = history_point0[best_slot] + rigid_contact_point1[i] = history_point1[best_slot] + else: + contact_penalty_k[i] = k_floor + contact_lambda[i] = wp.vec3(0.0) + + +@wp.kernel +def snapshot_contact_history( + rigid_contact_count: wp.array[int], + rigid_contact_shape0: wp.array[int], + rigid_contact_shape1: wp.array[int], + rigid_contact_point0: wp.array[wp.vec3], + rigid_contact_point1: wp.array[wp.vec3], + rigid_contact_normal: wp.array[wp.vec3], + contact_lambda: wp.array[wp.vec3], + contact_stick_flag: wp.array[wp.int32], + contact_penalty_k: wp.array[float], + # Hash table + ht_keys: wp.array[wp.uint64], + ht_active_slots: wp.array[wp.int32], + cell_size_inv: float, + # Outputs + history_point0: wp.array[wp.vec3], + history_point1: wp.array[wp.vec3], + history_lambda: wp.array[wp.vec3], + history_normal: wp.array[wp.vec3], + history_stick_flag: wp.array[wp.int32], + history_penalty_k: wp.array[float], + history_age: wp.array[int], + cached_slots: wp.array[int], +): + """Snapshot converged contact state into the history hash table (post-solve). + + Stores penalty_k, lambda, normal, stick state, and contact points. + In soft mode, lambda and stick_flag are inert but stored uniformly. + Caches the resolved slot index per contact so that subsequent + non-rebuild steps can update history without re-hashing. + """ + i = wp.tid() + if i >= rigid_contact_count[0]: + return + + s0 = rigid_contact_shape0[i] + s1 = rigid_contact_shape1[i] + p0 = rigid_contact_point0[i] + + pair_bits = _pack_shape_pair_bits(s0, s1) + cell = _compute_spatial_cell(p0, cell_size_inv) + key = _build_contact_history_key(pair_bits, cell) + + slot = hashtable_find_or_insert(key, ht_keys, ht_active_slots) + cached_slots[i] = slot + if slot >= 0: + history_point0[slot] = p0 + history_point1[slot] = rigid_contact_point1[i] + history_lambda[slot] = contact_lambda[i] + history_normal[slot] = rigid_contact_normal[i] + history_stick_flag[slot] = contact_stick_flag[i] + history_penalty_k[slot] = contact_penalty_k[i] + history_age[slot] = 0 @wp.kernel -def warmstart_body_particle_contacts( +def evict_contact_history( + ht_keys: wp.array[wp.uint64], + history_age: wp.array[int], + max_age: int, +): + """Increment age for all occupied history slots and evict stale entries. + + Scans the full hash table capacity (not active_slots) to avoid unbounded + active-list growth and ensure all persisted entries are properly aged. + + Note: snapshot_contact_history sets age = 0, and this kernel increments + before comparing, so effective retention is max_age + 1 steps. + """ + idx = wp.tid() + if ht_keys[idx] == HASHTABLE_EMPTY_KEY: + return + age = history_age[idx] + 1 + if age > max_age: + ht_keys[idx] = HASHTABLE_EMPTY_KEY + else: + history_age[idx] = age + + +@wp.kernel +def snapshot_contact_history_light( + rigid_contact_count: wp.array[int], + cached_slots: wp.array[int], + contact_lambda: wp.array[wp.vec3], + contact_stick_flag: wp.array[wp.int32], + contact_penalty_k: wp.array[float], + # Outputs (hash table history arrays, indexed by cached slot) + history_lambda: wp.array[wp.vec3], + history_stick_flag: wp.array[wp.int32], + history_penalty_k: wp.array[float], +): + """Write latest lambda/stick/penalty_k to the hash table at pre-cached slot positions. + + Runs on non-rebuild steps to keep the hash table fresh without hashing + or probing. Slots were cached by the previous full snapshot. + """ + i = wp.tid() + if i >= rigid_contact_count[0]: + return + slot = cached_slots[i] + if slot >= 0: + history_lambda[slot] = contact_lambda[i] + history_stick_flag[slot] = contact_stick_flag[i] + history_penalty_k[slot] = contact_penalty_k[i] + + +@wp.kernel +def step_body_body_contact_C0_lambda( + rigid_contact_count: wp.array[int], + rigid_contact_shape0: wp.array[int], + rigid_contact_shape1: wp.array[int], + rigid_contact_point0: wp.array[wp.vec3], + rigid_contact_point1: wp.array[wp.vec3], + rigid_contact_normal: wp.array[wp.vec3], + rigid_contact_margin0: wp.array[float], + rigid_contact_margin1: wp.array[float], + shape_body: wp.array[int], + body_q: wp.array[wp.transform], + hard_contacts: int, + gamma_lambda: float, + gamma_k: float, + contact_material_ke: wp.array[float], + k_start: float, + # In/out + contact_penalty_k: wp.array[float], + contact_C0: wp.array[wp.vec3], + contact_lambda: wp.array[wp.vec3], +): + """Per-step k decay + lambda decay + C0 snapshot. + + Runs every step. K decay is unconditional (hard and soft). Lambda decay + uses gamma_lambda (0 when contact history is off or contacts are soft). + C0 is always recomputed for hard contacts. + """ + i = wp.tid() + if i >= rigid_contact_count[0]: + return + + ke = contact_material_ke[i] + k_min = ke if k_start < 0.0 else wp.min(k_start, ke) + contact_penalty_k[i] = wp.clamp(gamma_k * contact_penalty_k[i], k_min, ke) + + contact_lambda[i] = contact_lambda[i] * gamma_lambda + + if hard_contacts == 1: + s0 = rigid_contact_shape0[i] + s1 = rigid_contact_shape1[i] + b0 = shape_body[s0] if s0 >= 0 else -1 + b1 = shape_body[s1] if s1 >= 0 else -1 + p0 = rigid_contact_point0[i] + p1 = rigid_contact_point1[i] + n = rigid_contact_normal[i] + cp0 = wp.transform_point(body_q[b0], p0) if b0 >= 0 else p0 + cp1 = wp.transform_point(body_q[b1], p1) if b1 >= 0 else p1 + thickness = rigid_contact_margin0[i] + rigid_contact_margin1[i] + d = cp1 - cp0 + contact_C0[i] = n * thickness - d + + +@wp.kernel +def init_body_particle_contacts( body_particle_contact_count: wp.array[int], body_particle_contact_shape: wp.array[int], soft_contact_ke: float, @@ -1902,42 +2444,40 @@ def warmstart_body_particle_contacts( shape_material_ke: wp.array[float], shape_material_kd: wp.array[float], shape_material_mu: wp.array[float], - k_start_body_contact: float, + k_start: float, # Outputs body_particle_contact_penalty_k: wp.array[float], - body_particle_contact_material_ke: wp.array[float], body_particle_contact_material_kd: wp.array[float], body_particle_contact_material_mu: wp.array[float], + body_particle_contact_material_ke: wp.array[float], ): - """ - Warm-start body-particle (particle-rigid) contact penalties and cache material - properties. - - The scalar inputs `soft_contact_ke/kd/mu` are the particle-side soft-contact - material parameters (from `model.soft_contact_*`). For each body-particle - contact, this kernel averages those particle-side values with the rigid - shape's material parameters and resets the AVBD penalty to a bounded - initial value based on `k_start_body_contact` and the averaged stiffness. + """Cold-start body-particle contact penalties and cache material properties. + + Averages particle-side material (scalar `soft_contact_ke/kd/mu`) with the + rigid shape's material. Penalty is seeded at ``min(k_start, avg_ke)`` when + ramping (k_start >= 0) or at ``avg_ke`` when fixed-k (k_start < 0). """ i = wp.tid() if i >= body_particle_contact_count[0]: return - # Read shape index for the rigid body side shape_idx = body_particle_contact_shape[i] - # Cache averaged material properties (arithmetic mean for stiffness/damping, geometric for friction) - avg_ke = 0.5 * (soft_contact_ke + shape_material_ke[shape_idx]) - avg_kd = 0.5 * (soft_contact_kd + shape_material_kd[shape_idx]) - avg_mu = wp.sqrt(soft_contact_mu * shape_material_mu[shape_idx]) + avg_ke, avg_kd, avg_mu = _average_contact_material( + soft_contact_ke, + soft_contact_kd, + soft_contact_mu, + shape_material_ke[shape_idx], + shape_material_kd[shape_idx], + shape_material_mu[shape_idx], + ) - body_particle_contact_material_ke[i] = avg_ke body_particle_contact_material_kd[i] = avg_kd body_particle_contact_material_mu[i] = avg_mu + body_particle_contact_material_ke[i] = avg_ke - # Reset contact penalty to k_start every frame because contact indices are not persistent across frames. - k_new = wp.min(k_start_body_contact, avg_ke) - body_particle_contact_penalty_k[i] = k_new + k_floor = avg_ke if k_start < 0.0 else wp.min(k_start, avg_ke) + body_particle_contact_penalty_k[i] = k_floor @wp.kernel @@ -1950,7 +2490,7 @@ def compute_cable_dahl_parameters( joint_X_p: wp.array[wp.transform], joint_X_c: wp.array[wp.transform], joint_constraint_start: wp.array[int], - joint_penalty_k_max: wp.array[float], + joint_penalty_k: wp.array[float], body_q: wp.array[wp.transform], body_q_rest: wp.array[wp.transform], joint_sigma_prev: wp.array[wp.vec3], @@ -2010,7 +2550,7 @@ def compute_cable_dahl_parameters( q_wc_rest = wp.transform_get_rotation(X_wc_rest) # Compute current curvature vector at beginning-of-step (predicted state) - kappa_now = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_now = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) # Read previous state (from last converged timestep) kappa_prev = joint_kappa_prev[j] @@ -2021,9 +2561,9 @@ def compute_cable_dahl_parameters( eps_max = joint_eps_max[j] tau = joint_tau[j] - # Use the per-joint bend stiffness target (cap) from the solver constraint caps (constraint slot 1 for cables). + # Use the per-joint bend stiffness from the solver constraint array (constraint slot 1 for cables). c_start = joint_constraint_start[j] - k_bend_target = joint_penalty_k_max[c_start + 1] + k_bend_target = joint_penalty_k[c_start + 1] # Friction envelope: sigma_max = k_bend_target * eps_max. @@ -2085,6 +2625,10 @@ def accumulate_body_body_contacts_per_body( contact_penalty_k: wp.array[float], contact_material_kd: wp.array[float], contact_material_mu: wp.array[float], + contact_lambda: wp.array[wp.vec3], + contact_C0: wp.array[wp.vec3], + avbd_alpha: float, + hard_contacts: int, rigid_contact_count: wp.array[int], rigid_contact_shape0: wp.array[int], rigid_contact_shape1: wp.array[int], @@ -2104,7 +2648,7 @@ def accumulate_body_body_contacts_per_body( body_hessian_aa: wp.array[wp.mat33], ): """ - Per-body contact accumulation with _NUM_CONTACT_THREADS_PER_BODY strided threads. + Per-body augmented-Lagrangian contact accumulation with _NUM_CONTACT_THREADS_PER_BODY strided threads. """ tid = wp.tid() body_idx_in_group = tid // _NUM_CONTACT_THREADS_PER_BODY @@ -2121,6 +2665,8 @@ def accumulate_body_body_contacts_per_body( if num_contacts > body_contact_buffer_pre_alloc: num_contacts = body_contact_buffer_pre_alloc + contact_count = rigid_contact_count[0] + force_acc = wp.vec3(0.0) torque_acc = wp.vec3(0.0) h_ll_acc = wp.mat33(0.0) @@ -2130,7 +2676,7 @@ def accumulate_body_body_contacts_per_body( i = thread_id_within_body while i < num_contacts: contact_idx = body_contact_indices[body_id * body_contact_buffer_pre_alloc + i] - if contact_idx >= rigid_contact_count[0]: + if contact_idx >= contact_count: i += _NUM_CONTACT_THREADS_PER_BODY continue @@ -2149,16 +2695,35 @@ def accumulate_body_body_contacts_per_body( cp0_world = wp.transform_point(body_q[b0], cp0_local) if b0 >= 0 else cp0_local cp1_world = wp.transform_point(body_q[b1], cp1_local) if b1 >= 0 else cp1_local thickness = rigid_contact_margin0[contact_idx] + rigid_contact_margin1[contact_idx] - dist = wp.dot(contact_normal, cp1_world - cp0_world) - penetration = thickness - dist + d = cp1_world - cp0_world + C_n = thickness - wp.dot(contact_normal, d) + + lam_n = float(0.0) + C_eff = C_n + lam_vec = wp.vec3(0.0) + k = contact_penalty_k[contact_idx] + + if hard_contacts == 1: + lam_vec = contact_lambda[contact_idx] + lam_n = wp.dot(lam_vec, contact_normal) + C0_vec = contact_C0[contact_idx] + C_vec = contact_normal * thickness - d + C_stab_vec = C_vec - avbd_alpha * C0_vec + C_stab_n = wp.dot(C_stab_vec, contact_normal) + C_eff = C_stab_n + + if C_n <= _SMALL_LENGTH_EPS and lam_n <= 0.0: + i += _NUM_CONTACT_THREADS_PER_BODY + continue - if penetration <= _SMALL_LENGTH_EPS: + f_n_check = k * C_eff + lam_n + if f_n_check <= 0.0 and lam_n <= 0.0: i += _NUM_CONTACT_THREADS_PER_BODY continue - contact_ke = contact_penalty_k[contact_idx] contact_kd = contact_material_kd[contact_idx] contact_mu = contact_material_mu[contact_idx] + ( force_0, torque_0, @@ -2179,11 +2744,14 @@ def accumulate_body_body_contacts_per_body( cp0_local, cp1_local, contact_normal, - penetration, - contact_ke, + C_eff, + k, + k, contact_kd, + lam_vec, contact_mu, friction_epsilon, + hard_contacts, dt, ) @@ -2219,8 +2787,8 @@ def compute_rigid_contact_forces( rigid_contact_point0: wp.array[wp.vec3], rigid_contact_point1: wp.array[wp.vec3], rigid_contact_normal: wp.array[wp.vec3], - rigid_contact_thickness0: wp.array[float], - rigid_contact_thickness1: wp.array[float], + rigid_contact_margin0: wp.array[float], + rigid_contact_margin1: wp.array[float], # Model/state shape_body: wp.array[wp.int32], body_q: wp.array[wp.transform], @@ -2230,6 +2798,10 @@ def compute_rigid_contact_forces( contact_penalty_k: wp.array[float], contact_material_kd: wp.array[float], contact_material_mu: wp.array[float], + contact_lambda: wp.array[wp.vec3], + contact_C0: wp.array[wp.vec3], + avbd_alpha: float, + hard_contacts: int, friction_epsilon: float, # Outputs (length = rigid_contact_max) out_body0: wp.array[wp.int32], @@ -2238,15 +2810,7 @@ def compute_rigid_contact_forces( out_point1_world: wp.array[wp.vec3], out_force_on_body1: wp.array[wp.vec3], ): - """Compute per-contact forces in world space, matching the AVBD rigid contact model. - - Produces a **contact-specific** force vector (world frame) for each rigid contact. - - Conventions: - - The computed force is the force **on body1** (shape1/body1) in world frame. - - The force on body0 is `-out_force_on_body1`. - - World-space contact points for both shapes are provided for wrench construction. - """ + """Compute per-contact forces in world space (hard: ALM, soft: penalty).""" contact_idx = wp.tid() rc = rigid_contact_count[0] @@ -2276,29 +2840,39 @@ def compute_rigid_contact_forces( cp0_local = rigid_contact_point0[contact_idx] cp1_local = rigid_contact_point1[contact_idx] - contact_normal = rigid_contact_normal[contact_idx] - # Static/kinematic shapes use shape_body == -1. In that case, contact points are already in world - # frame for that side and must not index into body_q[-1]. cp0_world = wp.transform_point(body_q[b0], cp0_local) if b0 >= 0 else cp0_local cp1_world = wp.transform_point(body_q[b1], cp1_local) if b1 >= 0 else cp1_local out_point0_world[contact_idx] = cp0_world out_point1_world[contact_idx] = cp1_world - thickness = rigid_contact_thickness0[contact_idx] + rigid_contact_thickness1[contact_idx] - dist = wp.dot(contact_normal, cp1_world - cp0_world) - penetration = thickness - dist - - if penetration <= _SMALL_LENGTH_EPS: + thickness = rigid_contact_margin0[contact_idx] + rigid_contact_margin1[contact_idx] + d = cp1_world - cp0_world + C_n = thickness - wp.dot(contact_normal, d) + + lam_n = float(0.0) + C_eff = C_n + lam_vec = wp.vec3(0.0) + k = contact_penalty_k[contact_idx] + + if hard_contacts == 1: + lam_vec = contact_lambda[contact_idx] + lam_n = wp.dot(lam_vec, contact_normal) + C0_vec = contact_C0[contact_idx] + C_vec = contact_normal * thickness - d + C_stab_vec = C_vec - avbd_alpha * C0_vec + C_stab_n = wp.dot(C_stab_vec, contact_normal) + C_eff = C_stab_n + + f_n_check = k * C_eff + lam_n + if (C_n <= _SMALL_LENGTH_EPS or f_n_check <= 0.0) and lam_n <= 0.0: out_force_on_body1[contact_idx] = wp.vec3(0.0) return - contact_ke = contact_penalty_k[contact_idx] contact_kd = contact_material_kd[contact_idx] contact_mu = contact_material_mu[contact_idx] - # Reuse the exact same contact model used by AVBD accumulation ( _force_0, _torque_0, @@ -2319,11 +2893,14 @@ def compute_rigid_contact_forces( cp0_local, cp1_local, contact_normal, - penetration, - contact_ke, + C_eff, + k, + k, contact_kd, + lam_vec, contact_mu, friction_epsilon, + hard_contacts, dt, ) @@ -2341,7 +2918,6 @@ def accumulate_body_particle_contacts_per_body( # Rigid body state body_q_prev: wp.array[wp.transform], body_q: wp.array[wp.transform], - body_qd: wp.array[wp.spatial_vector], body_com: wp.array[wp.vec3], body_inv_mass: wp.array[float], # AVBD body-particle soft contact penalties and material properties @@ -2352,13 +2928,9 @@ def accumulate_body_particle_contacts_per_body( # Soft contact data (body-particle) body_particle_contact_count: wp.array[int], body_particle_contact_particle: wp.array[int], - body_particle_contact_shape: wp.array[int], body_particle_contact_body_pos: wp.array[wp.vec3], body_particle_contact_body_vel: wp.array[wp.vec3], body_particle_contact_normal: wp.array[wp.vec3], - # Shape/material data - shape_material_mu: wp.array[float], - shape_body: wp.array[wp.int32], # Per-body soft-contact adjacency (body-particle) body_particle_contact_buffer_pre_alloc: int, body_particle_contact_counts: wp.array[wp.int32], @@ -2374,20 +2946,14 @@ def accumulate_body_particle_contacts_per_body( Per-body accumulation of body-particle (particle-rigid) soft contact forces and Hessians on rigid bodies. - This kernel mirrors the Gauss-Seidel per-body pattern used for body-body contacts, - but iterates over body-particle contacts associated with each body via the - precomputed body_particle_contact_* adjacency. - - For each body-particle contact, we: - 1. Reuse the particle-side contact model via evaluate_body_particle_contact to - compute the force and Hessian on the particle using adaptive AVBD penalties. - 2. Apply the equal-and-opposite reaction force, torque, and Hessian contributions - to the rigid body. + This kernel resolves contact geometry and relative displacement inline, then + calls ``_compute_body_particle_contact_force`` for the pure force law. + Body surface velocity uses the displacement-based path (body_q_prev). Notes: - Only dynamic bodies (inv_mass > 0) are updated. - Hessian contributions are accumulated into body_hessian_ll/al/aa. - - Uses AVBD adaptive penalty parameters (warmstarted and updated per iteration). + - Uses per-contact effective penalty/material parameters initialized once per step. """ tid = wp.tid() body_idx_in_group = tid // _NUM_CONTACT_THREADS_PER_BODY @@ -2404,9 +2970,19 @@ def accumulate_body_particle_contacts_per_body( if num_contacts > body_particle_contact_buffer_pre_alloc: num_contacts = body_particle_contact_buffer_pre_alloc - i = thread_id_within_body max_contacts = body_particle_contact_count[0] + X_wb = body_q[body_id] + X_wb_prev = body_q_prev[body_id] + com_world = wp.transform_point(X_wb, body_com[body_id]) + + force_acc = wp.vec3(0.0) + torque_acc = wp.vec3(0.0) + h_ll_acc = wp.mat33(0.0) + h_al_acc = wp.mat33(0.0) + h_aa_acc = wp.mat33(0.0) + + i = thread_id_within_body while i < num_contacts: contact_idx = body_particle_contact_indices[body_id * body_particle_contact_buffer_pre_alloc + i] if contact_idx >= max_contacts: @@ -2418,9 +2994,7 @@ def accumulate_body_particle_contacts_per_body( i += _NUM_CONTACT_THREADS_PER_BODY continue - # Early penetration check to avoid unnecessary function call particle_pos = particle_q[particle_idx] - X_wb = body_q[body_id] cp_local = body_particle_contact_body_pos[contact_idx] cp_world = wp.transform_point(X_wb, cp_local) n = body_particle_contact_normal[contact_idx] @@ -2431,64 +3005,44 @@ def accumulate_body_particle_contacts_per_body( i += _NUM_CONTACT_THREADS_PER_BODY continue - # Compute contact force and Hessian on particle using AVBD adaptive penalty - particle_prev_pos = particle_q_prev[particle_idx] - - # Read per-contact AVBD penalty and material properties - contact_ke = body_particle_contact_penalty_k[contact_idx] - contact_kd = body_particle_contact_material_kd[contact_idx] - contact_mu = body_particle_contact_material_mu[contact_idx] + bx_prev = wp.transform_point(X_wb_prev, cp_local) + bv = (cp_world - bx_prev) / dt + wp.transform_vector(X_wb, body_particle_contact_body_vel[contact_idx]) + dx = particle_pos - particle_q_prev[particle_idx] + relative_translation = dx - bv * dt - force_on_particle, hessian_particle = evaluate_body_particle_contact( - particle_idx, - particle_pos, - particle_prev_pos, - contact_idx, - contact_ke, - contact_kd, - contact_mu, + force_on_particle, hessian_particle = _compute_body_particle_contact_force( + penetration_depth, + n, + relative_translation, + body_particle_contact_penalty_k[contact_idx], + body_particle_contact_material_kd[contact_idx], + body_particle_contact_material_mu[contact_idx], friction_epsilon, - particle_radius, - shape_material_mu, - shape_body, - body_q, - body_q_prev, - body_qd, - body_com, - body_particle_contact_shape, - body_particle_contact_body_pos, - body_particle_contact_body_vel, - body_particle_contact_normal, dt, ) - # Reaction on the body (Newton's 3rd law) f_body = -force_on_particle - # Compute torque and Hessian contributions: tau = r x f, where r is from COM to contact point - com_world = wp.transform_point(X_wb, body_com[body_id]) r = cp_world - com_world tau_body = wp.cross(r, f_body) - # Hessian contributions for the rigid body - # The particle Hessian K_particle contributes to the body's linear-linear block - # and we need to add angular-linear and angular-angular blocks from the moment arm - K_total = hessian_particle r_skew = wp.skew(r) - r_skew_T_K = wp.transpose(r_skew) * K_total - - h_ll_body = K_total # Linear-linear block - h_al_body = -r_skew_T_K # Angular-linear block - h_aa_body = r_skew_T_K * r_skew # Angular-angular block + r_skew_T_K = wp.transpose(r_skew) * hessian_particle - wp.atomic_add(body_forces, body_id, f_body) - wp.atomic_add(body_torques, body_id, tau_body) - wp.atomic_add(body_hessian_ll, body_id, h_ll_body) - wp.atomic_add(body_hessian_al, body_id, h_al_body) - wp.atomic_add(body_hessian_aa, body_id, h_aa_body) + force_acc += f_body + torque_acc += tau_body + h_ll_acc += hessian_particle + h_al_acc += -r_skew_T_K + h_aa_acc += r_skew_T_K * r_skew i += _NUM_CONTACT_THREADS_PER_BODY + wp.atomic_add(body_forces, body_id, force_acc) + wp.atomic_add(body_torques, body_id, torque_acc) + wp.atomic_add(body_hessian_ll, body_id, h_ll_acc) + wp.atomic_add(body_hessian_al, body_id, h_al_acc) + wp.atomic_add(body_hessian_aa, body_id, h_aa_acc) + @wp.kernel def solve_rigid_body( @@ -2529,6 +3083,12 @@ def solve_rigid_body( joint_limit_upper: wp.array[float], joint_limit_ke: wp.array[float], joint_limit_kd: wp.array[float], + joint_lambda_lin: wp.array[wp.vec3], + joint_lambda_ang: wp.array[wp.vec3], + joint_C0_lin: wp.array[wp.vec3], + joint_C0_ang: wp.array[wp.vec3], + joint_is_hard: wp.array[wp.int32], + avbd_alpha: float, joint_dof_dim: wp.array2d[int], joint_rest_angle: wp.array[float], external_forces: wp.array[wp.vec3], @@ -2540,22 +3100,16 @@ def solve_rigid_body( body_q_new: wp.array[wp.transform], ): """ - AVBD solve step for rigid bodies using block Cholesky decomposition. - - Solves the 6-DOF rigid body system by assembling inertial, joint, and collision - contributions into a 6x6 block system: - - [ H_ll H_al^T ] - [ H_al H_aa ] + AVBD solve step for rigid bodies. - and solving via Schur complement. - Consistent with VBD particle solve pattern: inertia + external + constraint forces. + Assembles inertial, joint, and collision contributions into a 6x6 SPD + block system and solves via direct LDL^T. Algorithm: 1. Compute inertial forces/Hessians 2. Accumulate external forces/Hessians from rigid contacts 3. Accumulate joint forces/Hessians from adjacent joints - 4. Solve 6x6 block system via Schur complement: S = A - C*M^-1*C^T + 4. Solve 6x6 system via LDL^T 5. Update pose: rotation from angular increment, position from linear increment Args: @@ -2696,6 +3250,12 @@ def solve_rigid_body( joint_limit_upper, joint_limit_ke, joint_limit_kd, + joint_lambda_lin, + joint_lambda_ang, + joint_C0_lin, + joint_C0_ang, + joint_is_hard, + avbd_alpha, joint_dof_dim, joint_rest_angle, dt, @@ -2708,45 +3268,15 @@ def solve_rigid_body( h_al = h_al + joint_H_al h_aa = h_aa + joint_H_aa - # Solve 6x6 block system via Schur complement - # Regularize angular Hessian (in-place) + # Regularize angular Hessian trA = wp.trace(h_aa) / 3.0 epsA = 1.0e-9 * (trA + 1.0) h_aa[0, 0] = h_aa[0, 0] + epsA h_aa[1, 1] = h_aa[1, 1] + epsA h_aa[2, 2] = h_aa[2, 2] + epsA - # Factorize linear Hessian - Lm_p = chol33(h_ll) - - # Compute M^-1 * f_force - MinvF = chol33_solve(Lm_p, f_force) - - # Compute H_ll^{-1} * (H_al^T) - C_r0 = wp.vec3(h_al[0, 0], h_al[0, 1], h_al[0, 2]) - C_r1 = wp.vec3(h_al[1, 0], h_al[1, 1], h_al[1, 2]) - C_r2 = wp.vec3(h_al[2, 0], h_al[2, 1], h_al[2, 2]) - - X0 = chol33_solve(Lm_p, C_r0) - X1 = chol33_solve(Lm_p, C_r1) - X2 = chol33_solve(Lm_p, C_r2) - - # Columns are the solved vectors X0, X1, X2 - MinvCt = wp.mat33(X0[0], X1[0], X2[0], X0[1], X1[1], X2[1], X0[2], X1[2], X2[2]) - - # Compute Schur complement - S = h_aa - (h_al * MinvCt) - - # Factorize Schur complement - Ls_p = chol33(S) - - # Solve for angular increment - rhs_w = f_torque - (h_al * MinvF) - w_world = chol33_solve(Ls_p, rhs_w) - - # Solve for linear increment - Ct_w = wp.transpose(h_al) * w_world - x_inc = chol33_solve(Lm_p, f_force - Ct_w) + # Solve 6x6 system via direct LDL^T + x_inc, w_world = ldlt6_solve(h_ll, h_aa, h_al, f_force, f_torque) # Update pose from increments # Convert angular increment to quaternion @@ -2774,30 +3304,9 @@ def solve_rigid_body( body_q_new[body_index] = wp.transform(pos_new, rot_new) -@wp.kernel -def copy_rigid_body_transforms_back( - body_ids_in_color: wp.array[wp.int32], - body_q_in: wp.array[wp.transform], - body_q_out: wp.array[wp.transform], -): - """ - Copy body transforms for the current color group from source to destination. - - Used after the per-color solve to write updated world transforms from a temporary - buffer back to the main array. - - Args: - body_ids_in_color: Body indices in the current color group - body_q_in: Source body transforms (input) - body_q_out: Destination body transforms (output) - """ - tid = wp.tid() - body_index = body_ids_in_color[tid] - body_q_out[body_index] = body_q_in[body_index] - - @wp.kernel def update_duals_joint( + # Inputs joint_type: wp.array[int], joint_enabled: wp.array[bool], joint_parent: wp.array[int], @@ -2807,61 +3316,34 @@ def update_duals_joint( joint_axis: wp.array[wp.vec3], joint_qd_start: wp.array[int], joint_constraint_start: wp.array[int], - joint_penalty_k_max: wp.array[float], body_q: wp.array[wp.transform], body_q_rest: wp.array[wp.transform], - beta: float, - joint_penalty_k: wp.array[float], # input/output joint_dof_dim: wp.array2d[int], - joint_rest_angle: wp.array[float], - # Drive/limit parameters for adaptive drive/limit penalty growth + joint_C0_lin: wp.array[wp.vec3], + joint_C0_ang: wp.array[wp.vec3], + joint_is_hard: wp.array[wp.int32], + avbd_alpha: float, + joint_penalty_k_max: wp.array[float], + beta_lin: float, + beta_ang: float, joint_target_ke: wp.array[float], joint_target_pos: wp.array[float], joint_limit_lower: wp.array[float], joint_limit_upper: wp.array[float], joint_limit_ke: wp.array[float], + joint_rest_angle: wp.array[float], + # Input/output + joint_penalty_k: wp.array[float], + joint_lambda_lin: wp.array[wp.vec3], + joint_lambda_ang: wp.array[wp.vec3], ): """ - Update AVBD penalty parameters for joint constraints and drive/limits (per-iteration). - - Increases per-constraint penalties based on current constraint violation magnitudes, - clamped by the per-constraint stiffness cap ``joint_penalty_k_max``. - - For drive/limit slots, the cap is conditional on the active mode: - limit violated -> cap = ``model.joint_limit_ke``, drive active -> cap = ``model.joint_target_ke``. + Update augmented-Lagrangian duals for joint constraints (per-iteration). - Solver invariant: For each free DOF, drive and limit are mutually exclusive at runtime. - Limits have priority. The shared AVBD slot relies on this ordering. - - Notes: - - For ``JointType.CABLE``, ``joint_penalty_k_max`` is populated from ``model.joint_target_ke`` (material/constraint tuning). - - For rigid constraint slots like ``JointType.BALL``, ``JointType.FIXED``, ``JointType.REVOLUTE``, - ``JointType.PRISMATIC``, and ``JointType.D6``, ``joint_penalty_k_max`` is populated from solver-level caps - (e.g. ``rigid_joint_linear_ke`` / ``rigid_joint_angular_ke``). - - For drive/limit slots, ``joint_penalty_k_max`` stores the conservative warmstart ceiling - ``max(model.joint_target_ke, model.joint_limit_ke)``; the dual still clamps growth to the - active branch's cap (drive or limit) on each iteration. - - Args: - joint_type: Joint types - joint_parent: Parent body indices - joint_child: Child body indices - joint_X_p: Parent joint frames (local) - joint_X_c: Child joint frames (local) - joint_axis: Joint axis directions (per-DOF, used by REVOLUTE, PRISMATIC, and D6) - joint_qd_start: Joint DOF start indices (used to index into joint_axis) - joint_constraint_start: Start index per joint in the solver constraint layout - joint_penalty_k_max: Per-constraint stiffness cap (used for penalty clamping) - body_q: Current body transforms (world) - body_q_rest: Rest body transforms (world) - beta: Adaptation rate - joint_penalty_k: In/out per-constraint adaptive penalties - joint_dof_dim: Per-joint [lin_count, ang_count] for D6 joints - joint_target_ke: Model drive stiffness (for conditional cap in drive/limit dual update) - joint_target_pos: Control drive target position/angle (for drive violation computation) - joint_limit_lower: Model limit lower bound - joint_limit_upper: Model limit upper bound - joint_limit_ke: Model limit stiffness (for conditional cap in drive/limit dual update) + Structural slots (linear, angular) update lambda via ALM and ramp k, + both unconditionally. Drive/limit slots ramp k only (no lambda); + k is capped at ``joint_penalty_k_max`` while the force kernel applies + the mode-specific stiffness cap (``min(avbd_ke, model_ke)``). """ j = wp.tid() @@ -2899,50 +3381,67 @@ def update_duals_joint( X_wc = body_q[child] * joint_X_c[j] X_wc_rest = body_q_rest[child] * joint_X_c[j] - # Cable joint: adaptive penalty for stretch and bend constraints + # CABLE joint: isotropic stretch + isotropic bend penalties (2 scalars). if jt == JointType.CABLE: - # Read joint frame rotations q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) q_wp_rest = wp.transform_get_rotation(X_wp_rest) q_wc_rest = wp.transform_get_rotation(X_wc_rest) - # Compute scalar violation magnitudes used for penalty growth. - # Stretch uses meters residual magnitude ||x_c - x_p|| to update an effective [N/m] stiffness. x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) - C_stretch = wp.length(x_c - x_p) + C_vec_stretch = x_c - x_p - # Compute bend constraint violation (rotation vector magnitude) - kappa = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - C_bend = wp.length(kappa) + kappa = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - # Update stretch penalty (constraint slot 0) + # Stretch penalty update (constraint slot 0) stretch_idx = c_start - stiffness_stretch = joint_penalty_k_max[stretch_idx] - k_stretch = joint_penalty_k[stretch_idx] - k_stretch_new = wp.min(k_stretch + beta * C_stretch, stiffness_stretch) - joint_penalty_k[stretch_idx] = k_stretch_new + lam_new = _update_dual_vec3( + C_vec_stretch, + joint_C0_lin[j], + avbd_alpha, + joint_penalty_k[stretch_idx], + joint_lambda_lin[j], + joint_is_hard[stretch_idx], + ) + joint_lambda_lin[j] = lam_new + joint_penalty_k[stretch_idx] = wp.min( + joint_penalty_k_max[stretch_idx], joint_penalty_k[stretch_idx] + beta_lin * wp.length(C_vec_stretch) + ) - # Update bend penalty (constraint slot 1) + # Bend penalty update (constraint slot 1) bend_idx = c_start + 1 - stiffness_bend = joint_penalty_k_max[bend_idx] - k_bend = joint_penalty_k[bend_idx] - k_bend_new = wp.min(k_bend + beta * C_bend, stiffness_bend) - joint_penalty_k[bend_idx] = k_bend_new + lam_new = _update_dual_vec3( + kappa, + joint_C0_ang[j], + avbd_alpha, + joint_penalty_k[bend_idx], + joint_lambda_ang[j], + joint_is_hard[bend_idx], + ) + joint_lambda_ang[j] = lam_new + joint_penalty_k[bend_idx] = wp.min( + joint_penalty_k_max[bend_idx], joint_penalty_k[bend_idx] + beta_ang * wp.length(kappa) + ) return # BALL joint: update isotropic linear anchor-coincidence penalty (single scalar). if jt == JointType.BALL: - # Scalar violation magnitude used for penalty growth; force/Hessian uses C_vec directly. x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) C_vec = x_c - x_p - C_lin = wp.length(C_vec) i0 = c_start - k0 = joint_penalty_k[i0] - joint_penalty_k[i0] = wp.min(k0 + beta * C_lin, joint_penalty_k_max[i0]) + lam_new = _update_dual_vec3( + C_vec, + joint_C0_lin[j], + avbd_alpha, + joint_penalty_k[i0], + joint_lambda_lin[j], + joint_is_hard[i0], + ) + joint_lambda_lin[j] = lam_new + joint_penalty_k[i0] = wp.min(joint_penalty_k_max[i0], joint_penalty_k[i0] + beta_lin * wp.length(C_vec)) return # FIXED joint: update isotropic linear + isotropic angular penalties (2 scalars). @@ -2950,22 +3449,39 @@ def update_duals_joint( i_lin = c_start + 0 i_ang = c_start + 1 - # Linear violation magnitude x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) - C_lin = wp.length(x_c - x_p) - k_lin = joint_penalty_k[i_lin] - joint_penalty_k[i_lin] = wp.min(k_lin + beta * C_lin, joint_penalty_k_max[i_lin]) + C_vec_lin = x_c - x_p + lam_new = _update_dual_vec3( + C_vec_lin, + joint_C0_lin[j], + avbd_alpha, + joint_penalty_k[i_lin], + joint_lambda_lin[j], + joint_is_hard[i_lin], + ) + joint_lambda_lin[j] = lam_new + joint_penalty_k[i_lin] = wp.min( + joint_penalty_k_max[i_lin], joint_penalty_k[i_lin] + beta_lin * wp.length(C_vec_lin) + ) - # Angular violation magnitude from kappa q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) q_wp_rest = wp.transform_get_rotation(X_wp_rest) q_wc_rest = wp.transform_get_rotation(X_wc_rest) - kappa = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - C_ang = wp.length(kappa) - k_ang = joint_penalty_k[i_ang] - joint_penalty_k[i_ang] = wp.min(k_ang + beta * C_ang, joint_penalty_k_max[i_ang]) + kappa = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + lam_new = _update_dual_vec3( + kappa, + joint_C0_ang[j], + avbd_alpha, + joint_penalty_k[i_ang], + joint_lambda_ang[j], + joint_is_hard[i_ang], + ) + joint_lambda_ang[j] = lam_new + joint_penalty_k[i_ang] = wp.min( + joint_penalty_k_max[i_ang], joint_penalty_k[i_ang] + beta_ang * wp.length(kappa) + ) return # REVOLUTE joint: isotropic linear + perpendicular angular penalties (2 scalars). @@ -2975,23 +3491,40 @@ def update_duals_joint( qd_start = joint_qd_start[j] q_wp = wp.transform_get_rotation(X_wp) P_lin, P_ang = build_joint_projectors(jt, joint_axis, qd_start, 0, 1, q_wp) - a = wp.normalize(joint_axis[qd_start]) - # Linear violation magnitude (full, P_lin = I for REVOLUTE) x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) - C_lin = wp.length(P_lin * (x_c - x_p)) - k_lin = joint_penalty_k[i_lin] - joint_penalty_k[i_lin] = wp.min(k_lin + beta * C_lin, joint_penalty_k_max[i_lin]) + C_vec_lin = P_lin * (x_c - x_p) + lam_new = _update_dual_vec3( + C_vec_lin, + P_lin * joint_C0_lin[j], + avbd_alpha, + joint_penalty_k[i_lin], + joint_lambda_lin[j], + joint_is_hard[i_lin], + ) + joint_lambda_lin[j] = lam_new + joint_penalty_k[i_lin] = wp.min( + joint_penalty_k_max[i_lin], joint_penalty_k[i_lin] + beta_lin * wp.length(C_vec_lin) + ) - # Angular violation: projected perpendicular components of kappa q_wc = wp.transform_get_rotation(X_wc) q_wp_rest = wp.transform_get_rotation(X_wp_rest) q_wc_rest = wp.transform_get_rotation(X_wc_rest) - kappa = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - C_ang = wp.length(P_ang * kappa) - k_ang = joint_penalty_k[i_ang] - joint_penalty_k[i_ang] = wp.min(k_ang + beta * C_ang, joint_penalty_k_max[i_ang]) + kappa = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_perp = P_ang * kappa + lam_new = _update_dual_vec3( + kappa_perp, + P_ang * joint_C0_ang[j], + avbd_alpha, + joint_penalty_k[i_ang], + joint_lambda_ang[j], + joint_is_hard[i_ang], + ) + joint_lambda_ang[j] = lam_new + joint_penalty_k[i_ang] = wp.min( + joint_penalty_k_max[i_ang], joint_penalty_k[i_ang] + beta_ang * wp.length(kappa_perp) + ) # Drive/limit dual update for free angular DOF (slot c_start + 2) dof_idx = qd_start @@ -3003,19 +3536,16 @@ def update_duals_joint( has_limits = model_limit_ke > 0.0 and (lim_lower > -MAXVAL or lim_upper < MAXVAL) if has_drive or has_limits: + a = wp.normalize(joint_axis[qd_start]) theta = wp.dot(kappa, a) theta_abs = theta + joint_rest_angle[dof_idx] target_pos = joint_target_pos[dof_idx] - mode, err_pos = resolve_drive_limit_mode(theta_abs, target_pos, lim_lower, lim_upper, has_drive, has_limits) + _mode, err_pos = resolve_drive_limit_mode( + theta_abs, target_pos, lim_lower, lim_upper, has_drive, has_limits + ) i_dl = c_start + 2 - k_dl = joint_penalty_k[i_dl] C_dl = wp.abs(err_pos) - cap = joint_penalty_k_max[i_dl] - if mode == _DRIVE_LIMIT_MODE_LIMIT_LOWER or mode == _DRIVE_LIMIT_MODE_LIMIT_UPPER: - cap = model_limit_ke - elif mode == _DRIVE_LIMIT_MODE_DRIVE: - cap = model_drive_ke - joint_penalty_k[i_dl] = wp.min(k_dl + beta * C_dl, cap) + joint_penalty_k[i_dl] = wp.min(joint_penalty_k_max[i_dl], joint_penalty_k[i_dl] + beta_ang * C_dl) return # PRISMATIC joint: perpendicular linear + isotropic angular penalties (2 scalars). @@ -3023,26 +3553,43 @@ def update_duals_joint( i_lin = c_start + 0 i_ang = c_start + 1 qd_start = joint_qd_start[j] - axis_local = joint_axis[qd_start] q_wp = wp.transform_get_rotation(X_wp) P_lin, P_ang = build_joint_projectors(jt, joint_axis, qd_start, 1, 0, q_wp) - # Linear violation: projected perpendicular components x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) C_vec = x_c - x_p - C_lin = wp.length(P_lin * C_vec) - k_lin = joint_penalty_k[i_lin] - joint_penalty_k[i_lin] = wp.min(k_lin + beta * C_lin, joint_penalty_k_max[i_lin]) + C_vec_perp = P_lin * C_vec + lam_new = _update_dual_vec3( + C_vec_perp, + P_lin * joint_C0_lin[j], + avbd_alpha, + joint_penalty_k[i_lin], + joint_lambda_lin[j], + joint_is_hard[i_lin], + ) + joint_lambda_lin[j] = lam_new + joint_penalty_k[i_lin] = wp.min( + joint_penalty_k_max[i_lin], joint_penalty_k[i_lin] + beta_lin * wp.length(C_vec_perp) + ) - # Angular violation: full kappa (P_ang = I for PRISMATIC) q_wc = wp.transform_get_rotation(X_wc) q_wp_rest = wp.transform_get_rotation(X_wp_rest) q_wc_rest = wp.transform_get_rotation(X_wc_rest) - kappa = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - C_ang = wp.length(P_ang * kappa) - k_ang = joint_penalty_k[i_ang] - joint_penalty_k[i_ang] = wp.min(k_ang + beta * C_ang, joint_penalty_k_max[i_ang]) + kappa = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_perp = P_ang * kappa + lam_new = _update_dual_vec3( + kappa_perp, + P_ang * joint_C0_ang[j], + avbd_alpha, + joint_penalty_k[i_ang], + joint_lambda_ang[j], + joint_is_hard[i_ang], + ) + joint_lambda_ang[j] = lam_new + joint_penalty_k[i_ang] = wp.min( + joint_penalty_k_max[i_ang], joint_penalty_k[i_ang] + beta_ang * wp.length(kappa_perp) + ) # Drive/limit dual update for free linear DOF (slot c_start + 2) dof_idx = qd_start @@ -3054,19 +3601,14 @@ def update_duals_joint( has_limits = model_limit_ke > 0.0 and (lim_lower > -MAXVAL or lim_upper < MAXVAL) if has_drive or has_limits: + axis_local = joint_axis[qd_start] axis_w_dl = wp.normalize(wp.quat_rotate(q_wp, axis_local)) d_along = wp.dot(C_vec, axis_w_dl) target_pos = joint_target_pos[dof_idx] - mode, err_pos = resolve_drive_limit_mode(d_along, target_pos, lim_lower, lim_upper, has_drive, has_limits) + _mode, err_pos = resolve_drive_limit_mode(d_along, target_pos, lim_lower, lim_upper, has_drive, has_limits) i_dl = c_start + 2 - k_dl = joint_penalty_k[i_dl] C_dl = wp.abs(err_pos) - cap = joint_penalty_k_max[i_dl] - if mode == _DRIVE_LIMIT_MODE_LIMIT_LOWER or mode == _DRIVE_LIMIT_MODE_LIMIT_UPPER: - cap = model_limit_ke - elif mode == _DRIVE_LIMIT_MODE_DRIVE: - cap = model_drive_ke - joint_penalty_k[i_dl] = wp.min(k_dl + beta * C_dl, cap) + joint_penalty_k[i_dl] = wp.min(joint_penalty_k_max[i_dl], joint_penalty_k[i_dl] + beta_lin * C_dl) return # D6 joint: projected linear + projected angular penalties (2 scalars). @@ -3079,31 +3621,44 @@ def update_duals_joint( q_wp_rot = wp.transform_get_rotation(X_wp) P_lin, P_ang = build_joint_projectors(jt, joint_axis, qd_start, lin_count, ang_count, q_wp_rot) - # Linear violation: measure projected error x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) C_vec = x_c - x_p if lin_count < 3: - C_lin = wp.length(P_lin * C_vec) - else: - C_lin = 0.0 - k_lin = joint_penalty_k[i_lin] - joint_penalty_k[i_lin] = wp.min(k_lin + beta * C_lin, joint_penalty_k_max[i_lin]) + C_vec_perp = P_lin * C_vec + lam_new = _update_dual_vec3( + C_vec_perp, + P_lin * joint_C0_lin[j], + avbd_alpha, + joint_penalty_k[i_lin], + joint_lambda_lin[j], + joint_is_hard[i_lin], + ) + joint_lambda_lin[j] = lam_new + joint_penalty_k[i_lin] = wp.min( + joint_penalty_k_max[i_lin], joint_penalty_k[i_lin] + beta_lin * wp.length(C_vec_perp) + ) - # Angular violation: measure projected kappa q_wc = wp.transform_get_rotation(X_wc) q_wp_rest = wp.transform_get_rotation(X_wp_rest) q_wc_rest = wp.transform_get_rotation(X_wc_rest) - kappa = cable_get_kappa(q_wp_rot, q_wc, q_wp_rest, q_wc_rest) + kappa = compute_kappa(q_wp_rot, q_wc, q_wp_rest, q_wc_rest) if ang_count < 3: - C_ang = wp.length(P_ang * kappa) - else: - C_ang = 0.0 - k_ang = joint_penalty_k[i_ang] - joint_penalty_k[i_ang] = wp.min(k_ang + beta * C_ang, joint_penalty_k_max[i_ang]) + kappa_perp = P_ang * kappa + lam_new = _update_dual_vec3( + kappa_perp, + P_ang * joint_C0_ang[j], + avbd_alpha, + joint_penalty_k[i_ang], + joint_lambda_ang[j], + joint_is_hard[i_ang], + ) + joint_lambda_ang[j] = lam_new + joint_penalty_k[i_ang] = wp.min( + joint_penalty_k_max[i_ang], joint_penalty_k[i_ang] + beta_ang * wp.length(kappa_perp) + ) # Drive/limit dual update for D6 free DOFs - # Per free linear DOF for li in range(3): if li < lin_count: dof_idx = qd_start + li @@ -3118,20 +3673,13 @@ def update_duals_joint( axis_w_dl = wp.normalize(wp.quat_rotate(q_wp_rot, joint_axis[dof_idx])) d_along = wp.dot(C_vec, axis_w_dl) target_pos_dl = joint_target_pos[dof_idx] - mode, err_pos = resolve_drive_limit_mode( + _mode, err_pos = resolve_drive_limit_mode( d_along, target_pos_dl, lim_lower, lim_upper, has_drive, has_limits ) i_dl = c_start + 2 + li - k_dl = joint_penalty_k[i_dl] C_dl = wp.abs(err_pos) - cap = joint_penalty_k_max[i_dl] - if mode == _DRIVE_LIMIT_MODE_LIMIT_LOWER or mode == _DRIVE_LIMIT_MODE_LIMIT_UPPER: - cap = model_limit_ke - elif mode == _DRIVE_LIMIT_MODE_DRIVE: - cap = model_drive_ke - joint_penalty_k[i_dl] = wp.min(k_dl + beta * C_dl, cap) - - # Per free angular DOF + joint_penalty_k[i_dl] = wp.min(joint_penalty_k_max[i_dl], joint_penalty_k[i_dl] + beta_lin * C_dl) + for ai in range(3): if ai < ang_count: dof_idx = qd_start + lin_count + ai @@ -3147,18 +3695,12 @@ def update_duals_joint( theta = wp.dot(kappa, a_dl) theta_abs = theta + joint_rest_angle[dof_idx] target_pos_dl = joint_target_pos[dof_idx] - mode, err_pos = resolve_drive_limit_mode( + _mode, err_pos = resolve_drive_limit_mode( theta_abs, target_pos_dl, lim_lower, lim_upper, has_drive, has_limits ) i_dl = c_start + 2 + lin_count + ai - k_dl = joint_penalty_k[i_dl] C_dl = wp.abs(err_pos) - cap = joint_penalty_k_max[i_dl] - if mode == _DRIVE_LIMIT_MODE_LIMIT_LOWER or mode == _DRIVE_LIMIT_MODE_LIMIT_UPPER: - cap = model_limit_ke - elif mode == _DRIVE_LIMIT_MODE_DRIVE: - cap = model_drive_ke - joint_penalty_k[i_dl] = wp.min(k_dl + beta * C_dl, cap) + joint_penalty_k[i_dl] = wp.min(joint_penalty_k_max[i_dl], joint_penalty_k[i_dl] + beta_ang * C_dl) return @@ -3174,67 +3716,111 @@ def update_duals_body_body_contacts( rigid_contact_margin1: wp.array[float], shape_body: wp.array[int], body_q: wp.array[wp.transform], + body_q_prev: wp.array[wp.transform], + contact_material_mu: wp.array[float], + contact_C0: wp.array[wp.vec3], + avbd_alpha: float, + hard_contacts: int, + body_inv_mass: wp.array[float], contact_material_ke: wp.array[float], beta: float, - contact_penalty_k: wp.array[float], # input/output + # Input/output + contact_penalty_k: wp.array[float], + contact_lambda: wp.array[wp.vec3], + # Output + contact_stick_flag: wp.array[wp.int32], ): """ - Update AVBD penalty parameters for contact constraints (per-iteration). - - Increases each contact's penalty by beta * penetration and clamps to the - per-contact material stiffness. - - Args: - rigid_contact_count: Number of active contacts - rigid_contact_shape0/1: Shape ids for each contact pair - rigid_contact_point0/1: Contact points in local shape frames - rigid_contact_normal: Contact normals (pointing from shape0 to shape1) - rigid_contact_margin0/1: Per-shape margin (for SDF/capsule padding) - shape_body: Map from shape id to body id (-1 if kinematic/ground) - body_q: Current body transforms - contact_material_ke: Per-contact target stiffness - beta: Adaptation rate - contact_penalty_k: In/out per-contact adaptive penalties + Update AVBD augmented-Lagrangian duals for contact constraints (per-iteration). + Hard mode: scalar isotropic k with vec3 lambda. Normal uses C_stab_n, tangential + uses displacement (body_q_prev -> body_q) for kinematic friction support. + Coulomb cone clamping on tangential lambda. Soft mode: no lambda update. + K ramp runs unconditionally for both hard and soft contacts. """ idx = wp.tid() if idx >= rigid_contact_count[0]: return - # Read contact geometry shape_id_0 = rigid_contact_shape0[idx] shape_id_1 = rigid_contact_shape1[idx] body_id_0 = shape_body[shape_id_0] body_id_1 = shape_body[shape_id_1] - # Skip invalid contacts (both bodies kinematic/ground) if body_id_0 < 0 and body_id_1 < 0: return - # Read cached material stiffness - stiffness = contact_material_ke[idx] + cp0_local = rigid_contact_point0[idx] + cp1_local = rigid_contact_point1[idx] - # Transform contact points to world space if body_id_0 >= 0: - p0_world = wp.transform_point(body_q[body_id_0], rigid_contact_point0[idx]) + p0_world = wp.transform_point(body_q[body_id_0], cp0_local) + p0_prev = wp.transform_point(body_q_prev[body_id_0], cp0_local) else: - p0_world = rigid_contact_point0[idx] + p0_world = cp0_local + p0_prev = cp0_local if body_id_1 >= 0: - p1_world = wp.transform_point(body_q[body_id_1], rigid_contact_point1[idx]) + p1_world = wp.transform_point(body_q[body_id_1], cp1_local) + p1_prev = wp.transform_point(body_q_prev[body_id_1], cp1_local) else: - p1_world = rigid_contact_point1[idx] + p1_world = cp1_local + p1_prev = cp1_local - # Compute penetration depth (constraint violation) - # Distance along the A-to-B normal; positive implies separation + n = rigid_contact_normal[idx] d = p1_world - p0_world - dist = wp.dot(rigid_contact_normal[idx], d) thickness_total = rigid_contact_margin0[idx] + rigid_contact_margin1[idx] - penetration = wp.max(0.0, thickness_total - dist) - # Update penalty: k_new = min(k + beta * |C|, stiffness) - k = contact_penalty_k[idx] - k_new = wp.min(k + beta * penetration, stiffness) - contact_penalty_k[idx] = k_new + if hard_contacts == 1: + k = contact_penalty_k[idx] + C0_vec = contact_C0[idx] + lam_vec = contact_lambda[idx] + mu = contact_material_mu[idx] + + C_vec = n * thickness_total - d + C_stab_vec = C_vec - avbd_alpha * C0_vec + C_stab_n = wp.dot(C_stab_vec, n) + + # Release lambda_n at full rate on separation (bypass C0 stabilization). + C_n_raw = wp.dot(C_vec, n) + if C_n_raw < 0.0: + C_stab_n = C_n_raw + + lam_n_old = wp.dot(lam_vec, n) + + lam_n_new = wp.max(lam_n_old + k * C_stab_n, 0.0) + + rel_disp = (p0_world - p0_prev) - (p1_world - p1_prev) + rel_disp_n = wp.dot(n, rel_disp) + tangential_disp = rel_disp - n * rel_disp_n + lam_t_old = lam_vec - n * lam_n_old + lam_t_new = lam_t_old + k * tangential_disp + lam_t_len = wp.length(lam_t_new) + cone_limit = mu * lam_n_new + if lam_t_len > cone_limit and lam_t_len > 0.0: + lam_t_new = lam_t_new * (cone_limit / lam_t_len) + contact_lambda[idx] = n * lam_n_new + lam_t_new + + has_kinematic = int(0) + if body_id_0 < 0 or body_id_1 < 0: + has_kinematic = int(1) + elif body_id_0 >= 0 and body_inv_mass[body_id_0] == 0.0: + has_kinematic = int(1) + elif body_id_1 >= 0 and body_inv_mass[body_id_1] == 0.0: + has_kinematic = int(1) + + flag = int(0) + if lam_n_new > 0.0: + if has_kinematic == 1: + flag = _STICK_FLAG_ANCHOR + elif lam_t_len <= cone_limit: + flag = _STICK_FLAG_DEADZONE + contact_stick_flag[idx] = flag + else: + contact_stick_flag[idx] = int(0) + + C_n = wp.dot(n * thickness_total - d, n) + if C_n > 0.0: + contact_penalty_k[idx] = wp.min(contact_material_ke[idx], contact_penalty_k[idx] + beta * C_n) @wp.kernel @@ -3250,62 +3836,38 @@ def update_duals_body_particle_contacts( body_q: wp.array[wp.transform], body_particle_contact_material_ke: wp.array[float], beta: float, - body_particle_contact_penalty_k: wp.array[float], # input/output + body_particle_contact_penalty_k: wp.array[float], ): """ - Update AVBD penalty parameters for soft contact constraints (per-iteration). + Update AVBD penalty parameters for body-particle soft contacts (per-iteration). - Increases each soft contact's penalty by beta * penetration and clamps to the - per-contact material stiffness. - - Args: - body_particle_contact_count: Number of active body-particle soft contacts - body_particle_contact_particle: Particle index for each body-particle soft contact - body_particle_contact_shape: Shape index for each body-particle soft contact - body_particle_contact_body_pos: Contact points in local shape frames - body_particle_contact_normal: Contact normals (pointing from rigid to particle) - particle_q: Current particle positions - particle_radius: Particle radii - shape_body: Map from shape id to body id (-1 if kinematic/ground) - body_q: Current body transforms - body_particle_contact_material_ke: Per-contact target stiffness (averaged) - beta: Adaptation rate - body_particle_contact_penalty_k: In/out per-contact adaptive penalties + Ramps each contact's penalty by beta * penetration, clamped to the + per-contact material stiffness ceiling. """ idx = wp.tid() if idx >= body_particle_contact_count[0]: return - # Read contact data particle_idx = body_particle_contact_particle[idx] shape_idx = body_particle_contact_shape[idx] body_idx = shape_body[shape_idx] if shape_idx >= 0 else -1 - # Read cached material stiffness stiffness = body_particle_contact_material_ke[idx] - # Transform contact point to world space - # For rigid bodies (body_idx >= 0), transform from body-local to world space X_wb = wp.transform_identity() if body_idx >= 0: X_wb = body_q[body_idx] - cp_local = body_particle_contact_body_pos[idx] - cp_world = wp.transform_point(X_wb, cp_local) - - # Get particle data + cp_world = wp.transform_point(X_wb, body_particle_contact_body_pos[idx]) particle_pos = particle_q[particle_idx] radius = particle_radius[particle_idx] n = body_particle_contact_normal[idx] - # Compute penetration depth (constraint violation) penetration = -(wp.dot(n, particle_pos - cp_world) - radius) penetration = wp.max(0.0, penetration) - # Update penalty: k_new = min(k + beta * |C|, stiffness) k = body_particle_contact_penalty_k[idx] - k_new = wp.min(k + beta * penetration, stiffness) - body_particle_contact_penalty_k[idx] = k_new + body_particle_contact_penalty_k[idx] = wp.min(k + beta * penetration, stiffness) # ----------------------------- @@ -3316,23 +3878,50 @@ def update_body_velocity( dt: float, body_q: wp.array[wp.transform], body_com: wp.array[wp.vec3], + body_contact_buffer_pre_alloc: int, + body_contact_counts: wp.array[wp.int32], + body_contact_indices: wp.array[wp.int32], + contact_stick_flag: wp.array[wp.int32], + apply_stick_deadzone: int, + stick_freeze_translation_eps: float, + stick_freeze_angular_eps: float, body_q_prev: wp.array[wp.transform], body_qd: wp.array[wp.spatial_vector], + body_qd_mirror: wp.array[wp.spatial_vector], + body_q_out: wp.array[wp.transform], ): """ Update body velocities from position changes (world frame). + Optionally applies a tiny body-level stick-contact deadzone before + finite-difference velocity computation. Computes linear and angular velocities using finite differences. + Also transfers the final body poses to body_q_out (fused copy from + the in-place Gauss-Seidel iteration buffer to state_out). Linear: v = (com_current - com_prev) / dt Angular: omega from quaternion difference dq = q * q_prev^-1 Args: dt: Time step. - body_q: Current body transforms (world). + body_q: Current body transforms (world), from state_in (in-place iteration buffer). body_com: Center of mass offsets (local frame). - body_q_prev: Previous body transforms (input/output, advanced to current pose for next step). - body_qd: Output body velocities (spatial vectors, world frame). + body_contact_buffer_pre_alloc: Per-body contact-list capacity. + body_contact_counts: Number of body-body contacts adjacent to each body. + body_contact_indices: Flat per-body contact index lists. + contact_stick_flag: Per-contact flag (0=none, ANCHOR=kinematic, DEADZONE=dynamic-dynamic). + apply_stick_deadzone: If nonzero, enable anti-creep deadzone for bodies whose + contacts carry DEADZONE but not ANCHOR. + stick_freeze_translation_eps: Translation deadzone [m] for anti-creep snapping. + stick_freeze_angular_eps: Angular deadzone [rad] for anti-creep snapping. + body_q_prev: Previous body transforms (input/output, advanced to current + pose for next step). For kinematic bodies set body_q. For dynamic + teleportation also set body_q_prev and body_qd. + body_qd: Output body velocities (spatial vectors, world frame), bound to state_out. + body_qd_mirror: Output body velocities, bound to state_in. Mirrors body_qd so the + next step's forward integrator sees the finalized velocity even when the + caller's Python-level state swap is not recorded in a captured CUDA graph. + body_q_out: Output body transforms (state_out), fused copy of body_q. """ tid = wp.tid() @@ -3345,6 +3934,27 @@ def update_body_velocity( q = wp.transform_get_rotation(pose) q_prev = wp.transform_get_rotation(pose_prev) + if apply_stick_deadzone != 0: + count = wp.min(body_contact_counts[tid], body_contact_buffer_pre_alloc) + offset = tid * body_contact_buffer_pre_alloc + has_anchor = int(0) + has_deadzone = int(0) + for i in range(count): + contact_idx = body_contact_indices[offset + i] + f = contact_stick_flag[contact_idx] + if f == _STICK_FLAG_ANCHOR: + has_anchor = int(1) + elif f == _STICK_FLAG_DEADZONE: + has_deadzone = int(1) + + if has_deadzone != 0 and has_anchor == 0: + translation_delta = wp.length(x - x_prev) + angular_delta = wp.length(quat_velocity(q, q_prev, 1.0)) # dt=1 gives angular displacement [rad] + if translation_delta < stick_freeze_translation_eps and angular_delta < stick_freeze_angular_eps: + pose = pose_prev + x = x_prev + q = q_prev + # Compute COM positions com_local = body_com[tid] x_com = x + wp.quat_rotate(q, com_local) @@ -3358,9 +3968,14 @@ def update_body_velocity( body_qd[tid] = wp.spatial_vector(v, omega) + # Mirror to state_in (CUDA-graph-capture safety). + body_qd_mirror[tid] = wp.spatial_vector(v, omega) + # Advance body_q_prev for next step (for kinematic bodies this is the only write). body_q_prev[tid] = pose + body_q_out[tid] = pose + @wp.kernel def update_cable_dahl_state( @@ -3372,7 +3987,8 @@ def update_cable_dahl_state( joint_X_p: wp.array[wp.transform], joint_X_c: wp.array[wp.transform], joint_constraint_start: wp.array[int], - joint_penalty_k_max: wp.array[float], + joint_penalty_k: wp.array[float], + joint_is_hard: wp.array[wp.int32], # Body states (final, after solver convergence) body_q: wp.array[wp.transform], body_q_rest: wp.array[wp.transform], @@ -3395,7 +4011,7 @@ def update_cable_dahl_state( joint_parent, joint_child: Parent/child body indices joint_X_p, joint_X_c: Joint frames in parent/child joint_constraint_start: Start index per joint in the solver constraint layout - joint_penalty_k_max: Per-constraint stiffness cap; for cables, bend cap stores effective per-joint bend stiffness [N*m] + joint_penalty_k: Per-constraint penalty stiffness; for cables, bend slot stores effective per-joint bend stiffness [N*m] body_q: Final body transforms (after convergence) body_q_rest: Rest body transforms joint_sigma_prev: Friction stress state (read old, write new), wp.vec3 per joint @@ -3434,11 +4050,11 @@ def update_cable_dahl_state( q_wc_rest = wp.transform_get_rotation(X_wc_rest) # Compute final curvature vector at end of timestep - kappa_final = cable_get_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_final = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) - if not joint_enabled[j]: - # Refresh Dahl state to current configuration with zero preload so that - # re-enabling the joint does not see a stale kappa delta. + # Refresh Dahl state so toggling enabled/hard does not see stale values. + c_start_dahl = joint_constraint_start[j] + if not joint_enabled[j] or joint_is_hard[c_start_dahl + 1] == 1: joint_kappa_prev[j] = kappa_final joint_sigma_prev[j] = wp.vec3(0.0) joint_dkappa_prev[j] = wp.vec3(0.0) @@ -3453,9 +4069,9 @@ def update_cable_dahl_state( eps_max = joint_eps_max[j] # Maximum persistent strain [rad] tau = joint_tau[j] # Memory decay length [rad] - # Bend stiffness target (cap) is stored in constraint slot 1 for cable joints. + # Bend stiffness is stored in constraint slot 1 for cable joints. c_start = joint_constraint_start[j] - k_bend_target = joint_penalty_k_max[c_start + 1] # [N*m] + k_bend_target = joint_penalty_k[c_start + 1] # [N*m] # Friction envelope: sigma_max = k_bend_target * eps_max. sigma_max = k_bend_target * eps_max # [N*m] diff --git a/newton/_src/solvers/vbd/solver_vbd.py b/newton/_src/solvers/vbd/solver_vbd.py index db59f3fb2f..26fdc12b25 100644 --- a/newton/_src/solvers/vbd/solver_vbd.py +++ b/newton/_src/solvers/vbd/solver_vbd.py @@ -3,11 +3,14 @@ from __future__ import annotations +import warnings from typing import Any import numpy as np import warp as wp +from newton._src.geometry.hashtable import HashTable + from ...core.types import override from ...sim import ( Contacts, @@ -52,36 +55,36 @@ from .rigid_vbd_kernels import ( _NUM_CONTACT_THREADS_PER_BODY, RigidForceElementAdjacencyInfo, - # Adjacency building kernels _count_num_adjacent_joints, _fill_adjacent_joints, - # Iteration kernels - accumulate_body_body_contacts_per_body, # Body-body (rigid-rigid) contacts (Gauss-Seidel mode) - accumulate_body_particle_contacts_per_body, # Body-particle soft contacts (two-way coupling) - build_body_body_contact_lists, # Body-body (rigid-rigid) contact adjacency - build_body_particle_contact_lists, # Body-particle (rigid-particle) soft-contact adjacency - compute_cable_dahl_parameters, # Cable bending plasticity + accumulate_body_body_contacts_per_body, + accumulate_body_particle_contacts_per_body, + build_body_body_contact_lists, + build_body_particle_contact_lists, + check_contact_overflow, + compute_cable_dahl_parameters, compute_rigid_contact_forces, - copy_rigid_body_transforms_back, - # Pre-iteration kernels (rigid AVBD) + evict_contact_history, forward_step_rigid_bodies, + init_body_body_contact_materials, + init_body_body_contacts_avbd, + init_body_particle_contacts, + snapshot_contact_history, + snapshot_contact_history_light, solve_rigid_body, - # Post-iteration kernels + step_body_body_contact_C0_lambda, + step_joint_C0_lambda, update_body_velocity, update_cable_dahl_state, - update_duals_body_body_contacts, # Body-body (rigid-rigid) contacts (AVBD penalty update) - update_duals_body_particle_contacts, # Body-particle soft contacts (AVBD penalty update) - update_duals_joint, # Cable joints (AVBD penalty update) - warmstart_body_body_contacts, # Body-body (rigid-rigid) contacts (penalty warmstart) - warmstart_body_particle_contacts, # Body-particle soft contacts (penalty warmstart) - warmstart_joints, # Cable joints (stretch & bend) + update_duals_body_body_contacts, + update_duals_body_particle_contacts, + update_duals_joint, ) from .tri_mesh_collision import ( TriMeshCollisionDetector, TriMeshCollisionInfo, ) -# Export accumulate_contact_force_and_hessian for legacy collision_legacy.py compatibility __all__ = ["SolverVBD"] @@ -93,8 +96,12 @@ class SolverVBD(SolverBase): - Rigid body simulation (joints, contacts) using the AVBD algorithm - Coupled particle-rigid body systems - For rigid bodies, the AVBD algorithm uses **soft constraints** with adaptive penalty parameters - for joints and contacts. Hard constraints are not currently enforced. + For rigid bodies, the AVBD algorithm uses penalty stiffness (fixed by default, + optionally ramped via ``k_start``) together with augmented-Lagrangian state + for joints and contacts. + Rigid structural joint slots default to **hard mode** (augmented Lagrangian with persistent lambda + and C0 stabilization). Cable stretch defaults to **hard mode**, while cable bend defaults to + **soft mode**. The hard/soft mode can be changed per slot via :meth:`set_joint_constraint_mode`. Joint limitations: - Supported joint types: BALL, FIXED, FREE, REVOLUTE, PRISMATIC, D6, CABLE. @@ -118,6 +125,8 @@ class SolverVBD(SolverBase): References: - Anka He Chen, Ziheng Liu, Yin Yang, and Cem Yuksel. 2024. Vertex Block Descent. ACM Trans. Graph. 43, 4, Article 116 (July 2024), 16 pages. https://doi.org/10.1145/3658179 + - Chris Giles, Elie Diaz, and Cem Yuksel. 2025. Augmented Vertex Block Descent. ACM Trans. Graph. 44, 4, Article 90 (August 2025), 12 pages. + https://doi.org/10.1145/3731195 Note: `SolverVBD` requires coloring information for both particles and rigid bodies: @@ -156,6 +165,20 @@ class SolverVBD(SolverBase): state_in, state_out = state_out, state_in """ + class JointSlot: + """Named constraint slot indices for :meth:`set_joint_constraint_mode`. + + Structural slots (LINEAR=0, ANGULAR=1) map uniformly across all joint types: + - CABLE: LINEAR -> stretch, ANGULAR -> bend + - BALL: LINEAR only (1 slot) + - FIXED/REVOLUTE/PRISMATIC/D6: LINEAR + ANGULAR + """ + + LINEAR = 0 + ANGULAR = 1 + STRETCH = 0 + BEND = 1 + def __init__( self, model: Model, @@ -177,19 +200,30 @@ def __init__( particle_rest_shape_contact_exclusion_radius: float = 0.0, particle_external_vertex_contact_filtering_map: dict | None = None, particle_external_edge_contact_filtering_map: dict | None = None, - # Rigid body parameters - rigid_avbd_beta: float = 1.0e5, - rigid_avbd_gamma: float = 0.99, - rigid_contact_k_start: float = 1.0e2, # AVBD: initial stiffness for all body contacts (body-body + body-particle) - rigid_joint_linear_k_start: float = 1.0e4, # AVBD: initial stiffness seed for linear joint constraints - rigid_joint_angular_k_start: float = 1.0e1, # AVBD: initial stiffness seed for angular joint constraints - rigid_joint_linear_ke: float = 1.0e9, # AVBD: stiffness cap for non-cable linear joint constraints (BALL/FIXED/REVOLUTE/PRISMATIC/D6) - rigid_joint_angular_ke: float = 1.0e9, # AVBD: stiffness cap for non-cable angular joint constraints (FIXED/REVOLUTE/PRISMATIC/D6) - rigid_joint_linear_kd: float = 1.0e-2, # AVBD: Rayleigh damping coefficient for non-cable linear joint constraints - rigid_joint_angular_kd: float = 0.0, # AVBD: Rayleigh damping coefficient for non-cable angular joint constraints + # Rigid body parameters — AVBD hyperparameters + rigid_avbd_alpha: float = 0.95, # C0 stabilization strength (C_stab = C - alpha * C0) + rigid_avbd_joint_alpha: float | None = None, # Joint alpha override; None uses rigid_avbd_alpha + rigid_avbd_contact_alpha: float + | None = None, # Contact alpha override; None uses avbd_alpha (history+hard) or 0 + rigid_avbd_beta: float = 0.0, # Penalty ramp rate per iteration (0 = fixed-k) + rigid_avbd_linear_beta: float | None = None, # Linear beta override; None uses rigid_avbd_beta + rigid_avbd_angular_beta: float | None = None, # Angular beta override; None uses rigid_avbd_beta + rigid_avbd_gamma: float = 0.999, # Per-step decay for both k and lambda + # Rigid body - contacts + rigid_contact_hard: bool = True, # Hard = AL with lambda + C0 stabilization; soft = penalty only + rigid_contact_history: bool = False, # Persist contact state across steps (hard: k+lambda+C0; soft: k only) + rigid_contact_match_tolerance: float = 0.01, + rigid_contact_k_start: float = 1.0e2, # Contact penalty seed (used when beta > 0) rigid_body_contact_buffer_size: int = 64, rigid_body_particle_contact_buffer_size: int = 256, - rigid_enable_dahl_friction: bool = False, # Cable bending plasticity/hysteresis + # Rigid body - joints + rigid_joint_linear_ke: float = 1.0e5, # Penalty stiffness ceiling for structural linear joint constraints + rigid_joint_angular_ke: float = 1.0e5, # Penalty stiffness ceiling for structural angular joint constraints + rigid_joint_linear_k_start: float = 1.0e2, # Linear penalty seed (used when linear beta > 0) + rigid_joint_angular_k_start: float = 1.0e1, # Angular penalty seed (used when angular beta > 0) + rigid_joint_linear_kd: float = 0.0, # Rayleigh damping for non-cable linear joint constraints + rigid_joint_angular_kd: float = 0.0, # Rayleigh damping for non-cable angular joint constraints + rigid_enable_dahl_friction: bool | None = None, # Deprecated: auto-detected from model attributes ): """ Args: @@ -201,6 +235,8 @@ def __init__( iterations: Number of VBD iterations per step. friction_epsilon: Threshold to smooth small relative velocities in friction computation (used for both particle and rigid body contacts). + integrate_with_external_rigid_solver: Indicator for coupled rigid body-cloth simulation. When set to `True`, + the solver assumes rigid bodies are integrated by an external solver (one-way coupling). Particle parameters: @@ -210,8 +246,6 @@ def __init__( particle_self_contact_margin: The margin used for self-contact detection. This is the distance at which vertex-triangle pairs and edge-edge will be considered in contact generation. It should be larger than `particle_self_contact_radius` to avoid missing contacts. - integrate_with_external_rigid_solver: Indicator for coupled rigid body-cloth simulation. When set to `True`, - the solver assumes rigid bodies are integrated by an external solver (one-way coupling). particle_conservative_bound_relaxation: Relaxation factor for conservative penetration-free projection. particle_vertex_contact_buffer_size: Preallocation size for each vertex's vertex-triangle collision buffer. particle_edge_contact_buffer_size: Preallocation size for edge's edge-edge collision buffer. @@ -238,29 +272,45 @@ def __init__( Rigid body parameters: - rigid_avbd_beta: Penalty ramp rate for rigid body constraints (how fast k grows with constraint violation). - rigid_avbd_gamma: Warmstart decay for penalty k (cross-step decay factor for rigid body constraints). - rigid_contact_k_start: Initial penalty stiffness for all body contact constraints, including both body-body (rigid-rigid) - and body-particle (rigid-particle) contacts (AVBD). - rigid_joint_linear_k_start: Initial penalty seed for linear joint constraints (e.g., cable stretch, BALL linear). - Used to seed the per-constraint adaptive penalties for all linear joint constraints. - rigid_joint_angular_k_start: Initial penalty seed for angular joint constraints (e.g., cable bend, FIXED angular). - Used to seed the per-constraint adaptive penalties for all angular joint constraints. - rigid_joint_linear_ke: Stiffness cap used by AVBD for **non-cable** linear joint constraint scalars - (BALL, FIXED, REVOLUTE, PRISMATIC, and D6 projected linear slots). Cable joints use the - per-joint caps in ``model.joint_target_ke`` instead (cable interprets ``joint_target_ke/kd`` as - constraint tuning). - rigid_joint_angular_ke: Stiffness cap used by AVBD for **non-cable** angular joint constraint scalars - (FIXED, REVOLUTE, PRISMATIC, and D6 projected angular slots). - rigid_joint_linear_kd: Rayleigh damping coefficient for non-cable linear joint constraints (paired with - ``rigid_joint_linear_ke``). - rigid_joint_angular_kd: Rayleigh damping coefficient for non-cable angular joint constraints (paired with - ``rigid_joint_angular_ke``). - rigid_body_contact_buffer_size: Max body-body (rigid-rigid) contacts per rigid body for per-body contact lists (tune based on expected body-body contact density). - rigid_body_particle_contact_buffer_size: Max body-particle (rigid-particle) contacts per rigid body for per-body soft-contact lists (tune based on expected body-particle contact density). - rigid_enable_dahl_friction: Enable Dahl hysteresis friction model for cable bending (default: False). - Configure per-joint Dahl parameters via the solver-registered custom model attributes - ``model.vbd.dahl_eps_max`` and ``model.vbd.dahl_tau``. + rigid_avbd_alpha: C0 stabilization strength for all constraints (C_stab = C - alpha * C0). + Used as the default for both joints and contacts unless overridden. + rigid_avbd_joint_alpha: Joint-specific alpha override. ``None`` (default) + uses ``rigid_avbd_alpha``. + rigid_avbd_contact_alpha: Contact-specific alpha override. ``None`` (default) + uses ``rigid_avbd_alpha`` when history + hard contacts, ``0.0`` otherwise. + rigid_avbd_beta: Penalty ramp rate per AVBD iteration. ``0`` (default) disables + ramping (fixed-k). Set to e.g. ``1e5`` for ramping. Used for both linear and + angular constraints unless overridden. Note: linear (meters) and angular + (radians) constraints have different units, so the overrides should be used + for production tuning. + rigid_avbd_linear_beta: Linear beta override for linear constraints (meters). + ``None`` (default) uses ``rigid_avbd_beta``. + rigid_avbd_angular_beta: Angular beta override for angular constraints (radians). + ``None`` (default) uses ``rigid_avbd_beta``. + rigid_avbd_gamma: Per-step decay factor for both penalty k and lambda. Lower values + decay faster, improving stability at the cost of slower convergence. + rigid_contact_hard: Whether body-body rigid contacts use hard mode (augmented Lagrangian with + persistent lambda and C0 stabilization) or soft mode (penalty only). + rigid_contact_history: Whether to persist body-body contact state across steps via + spatial hash matching. For hard contacts, restores lambda, C0, and penalty k. + For soft contacts, restores penalty k only (useful with ramping). + rigid_contact_match_tolerance: World-space tolerance used to match current contacts against contact + history entries when ``rigid_contact_history=True``. + rigid_contact_k_start: Contact penalty seed for AVBD ramping. Used when ``rigid_avbd_beta > 0``. + When beta is 0, k is fixed at the contact stiffness regardless of this value. + rigid_body_contact_buffer_size: Max body-body contacts per rigid body for per-body contact lists. + rigid_body_particle_contact_buffer_size: Max body-particle contacts per rigid body for per-body + soft-contact lists. + rigid_joint_linear_ke: Penalty stiffness ceiling for non-cable structural linear joint slots. + rigid_joint_angular_ke: Penalty stiffness ceiling for non-cable structural angular joint slots. + rigid_joint_linear_k_start: Linear penalty seed for AVBD ramping. Used when linear beta > 0. + When beta is 0, k is fixed at the joint stiffness regardless of this value. + rigid_joint_angular_k_start: Angular penalty seed for AVBD ramping. Used when angular beta > 0. + When beta is 0, k is fixed at the joint stiffness regardless of this value. + rigid_joint_linear_kd: Rayleigh damping coefficient for non-cable linear joint constraints. + rigid_joint_angular_kd: Rayleigh damping coefficient for non-cable angular joint constraints. + rigid_enable_dahl_friction: Deprecated and ignored. Dahl friction is auto-detected + from ``model.vbd.dahl_eps_max`` / ``model.vbd.dahl_tau``. Note: - The `integrate_with_external_rigid_solver` argument enables one-way coupling between rigid body and soft body @@ -272,8 +322,23 @@ def __init__( Setting them too small may result in undetected collisions (particles) or contact overflow (rigid body contacts). Setting them excessively large may increase memory usage and degrade performance. + - Dahl hysteresis friction for cable bending is auto-detected from custom model attributes + ``model.vbd.dahl_eps_max`` and ``model.vbd.dahl_tau`` (set via ``register_custom_attributes``). """ + if rigid_enable_dahl_friction is not None: + warnings.warn( + "rigid_enable_dahl_friction is deprecated and ignored. " + "Dahl friction is now auto-detected from model attributes " + "(model.vbd.dahl_eps_max / model.vbd.dahl_tau). " + "To disable per-joint, set dahl_eps_max=0.", + DeprecationWarning, + stacklevel=2, + ) + + rigid_avbd_linear_beta = rigid_avbd_linear_beta if rigid_avbd_linear_beta is not None else rigid_avbd_beta + rigid_avbd_angular_beta = rigid_avbd_angular_beta if rigid_avbd_angular_beta is not None else rigid_avbd_beta + super().__init__(model) # Common parameters @@ -306,31 +371,30 @@ def __init__( # Initialize rigid body system and rigid-particle (body-particle) interaction state self._init_rigid_system( model, - rigid_avbd_beta, + rigid_avbd_alpha, + rigid_avbd_linear_beta, + rigid_avbd_angular_beta, rigid_avbd_gamma, + rigid_avbd_joint_alpha, + rigid_avbd_contact_alpha, + rigid_contact_hard, + rigid_contact_history, + rigid_contact_match_tolerance, rigid_contact_k_start, - rigid_joint_linear_k_start, - rigid_joint_angular_k_start, + rigid_body_contact_buffer_size, + rigid_body_particle_contact_buffer_size, rigid_joint_linear_ke, rigid_joint_angular_ke, + rigid_joint_linear_k_start, + rigid_joint_angular_k_start, rigid_joint_linear_kd, rigid_joint_angular_kd, - rigid_body_contact_buffer_size, - rigid_body_particle_contact_buffer_size, - rigid_enable_dahl_friction, ) - # Rigid-only flag to control whether to update cross-step history - # (rigid warmstart state such as contact/joint history). - # Defaults to True. This setting applies only to the next call to :meth:`step` and is then - # reset to ``True``. This is useful for substepping, where history update frequency might - # differ from the simulation step frequency (e.g. updating only on the first substep). - # This flag is automatically reset to True after each step(). - # Rigid warmstart update flag (contacts/joints). - self.update_rigid_history = True - - # Cached empty arrays for kernels that require wp.array arguments even when counts are zero. - self._empty_body_q = wp.empty(0, dtype=wp.transform, device=self.device) + # Controls whether the next step() refreshes rigid body-body contacts from + # the Contacts buffer or reuses the current rigid contact set. + # Defaults to True and auto-resets to True after each step(). + self._update_rigid_history = True def _init_particle_system( self, @@ -436,18 +500,24 @@ def _init_particle_system( def _init_rigid_system( self, model: Model, - rigid_avbd_beta: float, + rigid_avbd_alpha: float, + rigid_avbd_linear_beta: float, + rigid_avbd_angular_beta: float, rigid_avbd_gamma: float, + rigid_avbd_joint_alpha: float | None, + rigid_avbd_contact_alpha: float | None, + rigid_contact_hard: bool, + rigid_contact_history: bool, + rigid_contact_match_tolerance: float, rigid_contact_k_start: float, - rigid_joint_linear_k_start: float, - rigid_joint_angular_k_start: float, + rigid_body_contact_buffer_size: int, + rigid_body_particle_contact_buffer_size: int, rigid_joint_linear_ke: float, rigid_joint_angular_ke: float, + rigid_joint_linear_k_start: float, + rigid_joint_angular_k_start: float, rigid_joint_linear_kd: float, rigid_joint_angular_kd: float, - rigid_body_contact_buffer_size: int, - rigid_body_particle_contact_buffer_size: int, - rigid_enable_dahl_friction: bool, ): """Initialize rigid body-specific AVBD data structures and settings. @@ -456,15 +526,51 @@ def _init_rigid_system( - Shared interaction state for body-particle (rigid-particle) soft contacts """ # AVBD penalty parameters - self.avbd_beta = rigid_avbd_beta - self.avbd_gamma = rigid_avbd_gamma - - # Common initial penalty seed / lower bound for body contacts (clamped to non-negative) - self.k_start_body_contact = max(0.0, rigid_contact_k_start) - - # Joint constraint caps and damping for non-cable joints (constraint enforcement, not drives) - self.rigid_joint_linear_ke = max(0.0, rigid_joint_linear_ke) - self.rigid_joint_angular_ke = max(0.0, rigid_joint_angular_ke) + if rigid_avbd_linear_beta < 0: + raise ValueError(f"rigid_avbd_linear_beta must be >= 0, got {rigid_avbd_linear_beta}") + if rigid_avbd_angular_beta < 0: + raise ValueError(f"rigid_avbd_angular_beta must be >= 0, got {rigid_avbd_angular_beta}") + if not (0.0 <= rigid_avbd_gamma <= 1.0): + raise ValueError(f"rigid_avbd_gamma must be in [0, 1], got {rigid_avbd_gamma}") + if rigid_contact_k_start < 0: + raise ValueError(f"rigid_contact_k_start must be >= 0, got {rigid_contact_k_start}") + if rigid_joint_linear_k_start < 0: + raise ValueError(f"rigid_joint_linear_k_start must be >= 0, got {rigid_joint_linear_k_start}") + if rigid_joint_angular_k_start < 0: + raise ValueError(f"rigid_joint_angular_k_start must be >= 0, got {rigid_joint_angular_k_start}") + if rigid_joint_linear_ke < 0: + raise ValueError(f"rigid_joint_linear_ke must be >= 0, got {rigid_joint_linear_ke}") + if rigid_joint_angular_ke < 0: + raise ValueError(f"rigid_joint_angular_ke must be >= 0, got {rigid_joint_angular_ke}") + self.rigid_avbd_gamma = rigid_avbd_gamma + self.rigid_contact_k_start_value = -1.0 if rigid_avbd_linear_beta == 0.0 else float(rigid_contact_k_start) + self.rigid_joint_linear_k_start = rigid_joint_linear_k_start if rigid_avbd_linear_beta > 0.0 else None + self.rigid_joint_angular_k_start = rigid_joint_angular_k_start if rigid_avbd_angular_beta > 0.0 else None + # Resolve internal alpha (joint/contact) and beta (linear/angular) + self.rigid_joint_alpha = rigid_avbd_joint_alpha if rigid_avbd_joint_alpha is not None else rigid_avbd_alpha + self.rigid_linear_beta = rigid_avbd_linear_beta + self.rigid_angular_beta = rigid_avbd_angular_beta + self.rigid_contact_hard = int(rigid_contact_hard) + self.rigid_contact_history = rigid_contact_history + if rigid_avbd_contact_alpha is not None: + self.rigid_contact_alpha = rigid_avbd_contact_alpha + else: + # Zero alpha works better without contact history: no warm duals to stabilize. + self.rigid_contact_alpha = rigid_avbd_alpha if (rigid_contact_history and rigid_contact_hard) else 0.0 + rigid_contact_match_tolerance = max(0.0, rigid_contact_match_tolerance) + self.rigid_contact_match_tolerance = rigid_contact_match_tolerance + self.rigid_contact_match_cell_size_inv = 1.0 / max(rigid_contact_match_tolerance, 1.0e-8) + self.rigid_contact_match_tolerance_sq = rigid_contact_match_tolerance * rigid_contact_match_tolerance + + # DEADZONE body-snap thresholds; suppressed by _STICK_FLAG_ANCHOR. + self.rigid_contact_stick_freeze_translation_eps = 1.0e-4 + self.rigid_contact_stick_freeze_angular_eps = 1.0e-4 + + self._rigid_contact_match_max_age = 0 # multi-frame persistence untested; clear+refresh each frame + + # Joint constraint stiffness and damping for non-cable structural joints + self.rigid_joint_linear_ke = rigid_joint_linear_ke + self.rigid_joint_angular_ke = rigid_joint_angular_ke self.rigid_joint_linear_kd = max(0.0, rigid_joint_linear_kd) self.rigid_joint_angular_kd = max(0.0, rigid_joint_angular_kd) @@ -472,9 +578,10 @@ def _init_rigid_system( # Rigid-only AVBD state (used when SolverVBD integrates bodies) # ------------------------------------------------------------- if not self.integrate_with_external_rigid_solver and model.body_count > 0: - # State storage - # Initialize to the current poses for the first step to avoid spurious finite-difference - # velocities/friction impulses. + # Previous-step body transforms, advanced by update_body_velocity() each step. + # Provides contact friction velocity and joint C0 feedforward for kinematic tracking. + # Kinematic bodies: set body_q. + # Dynamic teleportation: also set body_q_prev and body_qd. self.body_q_prev = wp.clone(model.body_q).to(self.device) self.body_inertia_q = wp.zeros_like(model.body_q, device=self.device) # inertial target poses for AVBD @@ -485,53 +592,64 @@ def _init_rigid_system( self.body_torques = wp.zeros(self.model.body_count, dtype=wp.vec3, device=self.device) self.body_forces = wp.zeros(self.model.body_count, dtype=wp.vec3, device=self.device) + # Persistent scratch for joint_f accumulation + self._body_f_for_integration = wp.zeros(self.model.body_count, dtype=wp.spatial_vector, device=self.device) + # Hessian blocks (6x6 block structure: angular-angular, angular-linear, linear-linear) self.body_hessian_aa = wp.zeros(self.model.body_count, dtype=wp.mat33, device=self.device) self.body_hessian_al = wp.zeros(self.model.body_count, dtype=wp.mat33, device=self.device) self.body_hessian_ll = wp.zeros(self.model.body_count, dtype=wp.mat33, device=self.device) - # Per-body contact lists - # Body-body (rigid-rigid) contact adjacency (CSR-like: per-body counts and flat index array) - self.body_body_contact_buffer_pre_alloc = rigid_body_contact_buffer_size + # Per-body contact lists (CSR-like: per-body counts + flat index array). + # Tight: pre_alloc = 0 when the contact source is absent (no shapes / no particles). + bb_pre_alloc = rigid_body_contact_buffer_size if model.shape_count > 0 else 0 + self.body_body_contact_buffer_pre_alloc = bb_pre_alloc self.body_body_contact_counts = wp.zeros(self.model.body_count, dtype=wp.int32, device=self.device) self.body_body_contact_indices = wp.zeros( - self.model.body_count * self.body_body_contact_buffer_pre_alloc, dtype=wp.int32, device=self.device + self.model.body_count * bb_pre_alloc, dtype=wp.int32, device=self.device ) + self.body_body_contact_overflow_max = wp.zeros(1, dtype=wp.int32, device=self.device) - # Body-particle (rigid-particle) contact adjacency (CSR-like: per-body counts and flat index array) - self.body_particle_contact_buffer_pre_alloc = rigid_body_particle_contact_buffer_size + bp_pre_alloc = rigid_body_particle_contact_buffer_size if model.particle_count > 0 else 0 + self.body_particle_contact_buffer_pre_alloc = bp_pre_alloc self.body_particle_contact_counts = wp.zeros(self.model.body_count, dtype=wp.int32, device=self.device) self.body_particle_contact_indices = wp.zeros( - self.model.body_count * self.body_particle_contact_buffer_pre_alloc, - dtype=wp.int32, - device=self.device, + self.model.body_count * bp_pre_alloc, dtype=wp.int32, device=self.device ) + self.body_particle_contact_overflow_max = wp.zeros(1, dtype=wp.int32, device=self.device) - # AVBD constraint penalties - # Joint constraint layout + penalties (solver constraint scalars) + # Joint constraint layout + penalty stiffness (mutable k, frozen bounds) self._init_joint_constraint_layout() - self.joint_penalty_k = self._init_joint_penalty_k(rigid_joint_linear_k_start, rigid_joint_angular_k_start) + self.joint_penalty_k, self.joint_penalty_k_min, self.joint_penalty_k_max = self._init_joint_penalty_k() self.joint_rest_angle = self._init_joint_rest_angle() - # Contact penalties (adaptive penalties for body-body contacts) - if model.shape_count > 0: - max_contacts = getattr(model, "rigid_contact_max", 0) or 0 - if max_contacts <= 0: - # Estimate from shape contact pairs (same heuristic previously in finalize()) - pair_count = model.shape_contact_pair_count if hasattr(model, "shape_contact_pair_count") else 0 - max_contacts = max(10000, pair_count * 20) - # Per-contact AVBD penalty for body-body contacts - self.body_body_contact_penalty_k = wp.full( - (max_contacts,), self.k_start_body_contact, dtype=float, device=self.device - ) - - # Pre-computed averaged body-body contact material properties (computed once per step in warmstart) - self.body_body_contact_material_ke = wp.zeros(max_contacts, dtype=float, device=self.device) - self.body_body_contact_material_kd = wp.zeros(max_contacts, dtype=float, device=self.device) - self.body_body_contact_material_mu = wp.zeros(max_contacts, dtype=float, device=self.device) - - # Dahl friction model (cable bending plasticity) - # State variables for Dahl hysteresis (persistent across timesteps) + # Body-body contact state (pre-allocated in __init__ when possible, resized on first step otherwise). + self.body_body_contact_penalty_k = wp.zeros(0, dtype=float, device=self.device) + self.body_body_contact_material_ke = wp.zeros(0, dtype=float, device=self.device) + self.body_body_contact_material_kd = wp.zeros(0, dtype=float, device=self.device) + self.body_body_contact_material_mu = wp.zeros(0, dtype=float, device=self.device) + self.body_body_contact_lambda = wp.zeros(0, dtype=wp.vec3, device=self.device) + self.body_body_contact_C0 = wp.zeros(0, dtype=wp.vec3, device=self.device) + self.body_body_contact_stick_flag = wp.zeros(0, dtype=wp.int32, device=self.device) + + # Contact history (same pattern as body-body contact state). + self.contact_history_ht = None + self.contact_history_point0 = None + self.contact_history_point1 = None + self.contact_history_lambda = None + self.contact_history_normal = None + self.contact_history_stick_flag = None + self.contact_history_penalty_k = None + self.contact_history_age = None + self.contact_history_slots = None + + # Joint augmented-Lagrangian state (vec3, per-joint, bilateral) + self.joint_lambda_lin = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) + self.joint_lambda_ang = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) + self.joint_C0_lin = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) + self.joint_C0_ang = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) + + # Dahl friction state (cable bending plasticity, persistent across timesteps) self.joint_sigma_prev = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) self.joint_kappa_prev = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) self.joint_dkappa_prev = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) @@ -540,51 +658,55 @@ def _init_rigid_system( self.joint_sigma_start = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) self.joint_C_fric = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) - # Dahl model configuration - self.enable_dahl_friction = rigid_enable_dahl_friction - self.joint_dahl_eps_max = wp.zeros(model.joint_count, dtype=float, device=self.device) - self.joint_dahl_tau = wp.zeros(model.joint_count, dtype=float, device=self.device) - - if rigid_enable_dahl_friction: - if model.joint_count == 0: - self.enable_dahl_friction = False - else: - # Read per-joint Dahl parameters from model.vbd if present; otherwise use defaults (eps_max=0.5, tau=1.0). - # Recommended: call SolverVBD.register_custom_attributes(builder) before finalize() to allocate these arrays. - vbd_attrs: Any = getattr(model, "vbd", None) - if vbd_attrs is not None and hasattr(vbd_attrs, "dahl_eps_max") and hasattr(vbd_attrs, "dahl_tau"): - self.joint_dahl_eps_max = vbd_attrs.dahl_eps_max - self.joint_dahl_tau = vbd_attrs.dahl_tau - else: - self._init_dahl_params(0.5, 1.0, model) + # Dahl friction: auto-detected from model.vbd custom attributes. + # Enabled when model has dahl_eps_max and dahl_tau (set via register_custom_attributes). + vbd_attrs: Any = getattr(model, "vbd", None) + has_dahl = ( + model.joint_count > 0 + and vbd_attrs is not None + and hasattr(vbd_attrs, "dahl_eps_max") + and hasattr(vbd_attrs, "dahl_tau") + ) + self.enable_dahl_friction = has_dahl + if has_dahl: + self.joint_dahl_eps_max = vbd_attrs.dahl_eps_max + self.joint_dahl_tau = vbd_attrs.dahl_tau + else: + self.joint_dahl_eps_max = wp.zeros(model.joint_count, dtype=float, device=self.device) + self.joint_dahl_tau = wp.zeros(model.joint_count, dtype=float, device=self.device) # ------------------------------------------------------------- - # Body-particle interaction - shared state + # Body-particle interaction - shared state (pre-allocated below when + # particles exist, otherwise empty). # ------------------------------------------------------------- - # Soft contact penalties (adaptive penalties for body-particle contacts) - # Use same initial penalty as body-body contacts - max_soft_contacts = model.shape_count * model.particle_count - # Per-contact AVBD penalty for body-particle soft contacts (same initial seed as body-body) - self.body_particle_contact_penalty_k = wp.full( - (max_soft_contacts,), self.k_start_body_contact, dtype=float, device=self.device - ) - - # Pre-computed averaged body-particle soft contact material properties (computed once per step in warmstart) - # These correspond to body-particle soft contacts and are averaged between model.soft_contact_* - # and shape material properties. - self.body_particle_contact_material_ke = wp.zeros(max_soft_contacts, dtype=float, device=self.device) - self.body_particle_contact_material_kd = wp.zeros(max_soft_contacts, dtype=float, device=self.device) - self.body_particle_contact_material_mu = wp.zeros(max_soft_contacts, dtype=float, device=self.device) + self.body_particle_contact_penalty_k = wp.zeros(0, dtype=float, device=self.device) + self.body_particle_contact_material_ke = wp.zeros(0, dtype=float, device=self.device) + self.body_particle_contact_material_kd = wp.zeros(0, dtype=float, device=self.device) + self.body_particle_contact_material_mu = wp.zeros(0, dtype=float, device=self.device) # Kinematic body support: create effective inv_mass / inv_inertia arrays # with kinematic bodies zeroed out. self._init_kinematic_state() + # Pre-allocate body-body contact buffers when the collision pipeline + # has already been created (model.rigid_contact_max > 0). When the + # solver is created before model.contacts(), rigid_contact_max is 0 + # and we fall back to lazy allocation in _initialize_rigid_bodies. + rcm = getattr(self.model, "rigid_contact_max", 0) or 0 + if rcm > 0 and self.model.body_count > 0 and not self.integrate_with_external_rigid_solver: + self._init_body_body_contact_state(rcm) + if self.rigid_contact_history: + self._init_contact_history(rcm) + # Body-particle buffers are always pre-allocated from model topology + # (shape_count * particle_count); no pipeline dependency. + if self.model.particle_count > 0 and self.model.shape_count > 0: + self._init_body_particle_contact_state(self.model.shape_count * self.model.particle_count) + # Validation has_bodies = self.model.body_count > 0 has_body_coloring = len(self.model.body_color_groups) > 0 - if has_bodies and not has_body_coloring: + if has_bodies and not has_body_coloring and not self.integrate_with_external_rigid_solver: raise ValueError( "model.body_color_groups is empty but rigid bodies are present! When using the SolverVBD you must call ModelBuilder.color() " "or ModelBuilder.set_coloring() before calling ModelBuilder.finalize()." @@ -599,36 +721,64 @@ def notify_model_changed(self, flags: int) -> None: # Initialization Helper Methods # ===================================================== + def _init_body_body_contact_state(self, rigid_contact_max: int) -> None: + """Allocate body-body contact state arrays sized to the given contact buffer capacity.""" + self.body_body_contact_penalty_k = wp.zeros(rigid_contact_max, dtype=float, device=self.device) + self.body_body_contact_material_ke = wp.zeros(rigid_contact_max, dtype=float, device=self.device) + self.body_body_contact_material_kd = wp.zeros(rigid_contact_max, dtype=float, device=self.device) + self.body_body_contact_material_mu = wp.zeros(rigid_contact_max, dtype=float, device=self.device) + self.body_body_contact_lambda = wp.zeros(rigid_contact_max, dtype=wp.vec3, device=self.device) + self.body_body_contact_C0 = wp.zeros(rigid_contact_max, dtype=wp.vec3, device=self.device) + self.body_body_contact_stick_flag = wp.zeros(rigid_contact_max, dtype=wp.int32, device=self.device) + + def _init_body_particle_contact_state(self, soft_contact_max: int) -> None: + """Allocate body-particle material arrays sized to the given soft contact capacity.""" + self.body_particle_contact_penalty_k = wp.zeros(soft_contact_max, dtype=float, device=self.device) + self.body_particle_contact_material_ke = wp.zeros(soft_contact_max, dtype=float, device=self.device) + self.body_particle_contact_material_kd = wp.zeros(soft_contact_max, dtype=float, device=self.device) + self.body_particle_contact_material_mu = wp.zeros(soft_contact_max, dtype=float, device=self.device) + + def _init_contact_history(self, rigid_contact_max: int) -> None: + """Allocate the contact history hash table sized to the given contact buffer capacity.""" + ht_capacity = max(256, 4 * rigid_contact_max) + self.contact_history_ht = HashTable(capacity=ht_capacity, device=self.device) + ht_cap = self.contact_history_ht.capacity + self.contact_history_point0 = wp.zeros(ht_cap, dtype=wp.vec3, device=self.device) + self.contact_history_point1 = wp.zeros(ht_cap, dtype=wp.vec3, device=self.device) + self.contact_history_lambda = wp.zeros(ht_cap, dtype=wp.vec3, device=self.device) + self.contact_history_normal = wp.zeros(ht_cap, dtype=wp.vec3, device=self.device) + self.contact_history_stick_flag = wp.zeros(ht_cap, dtype=wp.int32, device=self.device) + self.contact_history_penalty_k = wp.zeros(ht_cap, dtype=float, device=self.device) + self.contact_history_age = wp.zeros(ht_cap, dtype=int, device=self.device) + self.contact_history_slots = wp.full(rigid_contact_max, -1, dtype=int, device=self.device) + + @staticmethod + def _to_numpy(arr, dtype=None): + """Transfer a Warp array to CPU and return as numpy, optionally casting dtype.""" + cpu = arr.to("cpu") + result = cpu.numpy() if hasattr(cpu, "numpy") else np.asarray(cpu) + return result if dtype is None else result.astype(dtype, copy=False) + def _init_joint_constraint_layout(self) -> None: """Initialize VBD-owned joint constraint indexing. - VBD stores and adapts penalty stiffness values for *scalar constraint components*: - - ``JointType.CABLE``: 2 scalars (stretch/linear, bend/angular) - - ``JointType.BALL``: 1 scalar (isotropic linear anchor-coincidence) - - ``JointType.FIXED``: 2 scalars (isotropic linear anchor-coincidence + isotropic angular) - - ``JointType.REVOLUTE``: 3 scalars (isotropic linear + 2-DOF perpendicular angular + angular drive/limit) - - ``JointType.PRISMATIC``: 3 scalars (2-DOF perpendicular linear + isotropic angular + linear drive/limit) - - ``JointType.D6``: 2 + lin_count + ang_count scalars (projected linear + projected angular + per-DOF drive/limit) - - ``JointType.FREE``: 0 scalars (not a constraint) - - Ordering (must match kernel indexing via ``joint_constraint_start``): - - ``JointType.CABLE``: [stretch (linear), bend (angular)] - - ``JointType.BALL``: [linear] - - ``JointType.FIXED``: [linear, angular] - - ``JointType.REVOLUTE``: [linear, angular, ang_drive_limit] - - ``JointType.PRISMATIC``: [linear, angular, lin_drive_limit] - - ``JointType.D6``: [linear, angular, lin_dl_0, ..., ang_dl_0, ...] + VBD stores and adapts penalty stiffness values for scalar constraint components: + - CABLE: 2 scalars (stretch/linear, bend/angular) + - BALL: 1 scalar (isotropic linear anchor-coincidence) + - FIXED: 2 scalars (isotropic linear + isotropic angular) + - REVOLUTE: 3 scalars (isotropic linear + 2-DOF perpendicular angular + angular drive/limit) + - PRISMATIC: 3 scalars (2-DOF perpendicular linear + isotropic angular + linear drive/limit) + - D6: 2 + lin_count + ang_count scalars (projected linear + projected angular + per-DOF drive/limit) + - FREE: 0 scalars (not a constraint) Drive and limit for each free DOF share one AVBD slot (mutually exclusive at runtime). - Any other joint type will raise ``NotImplementedError``. + Any other joint type will raise NotImplementedError. """ n_j = self.model.joint_count with wp.ScopedDevice("cpu"): - jt_cpu = self.model.joint_type.to("cpu") - jt = jt_cpu.numpy() if hasattr(jt_cpu, "numpy") else np.asarray(jt_cpu, dtype=int) - jdof_dim_cpu = self.model.joint_dof_dim.to("cpu") - jdof_dim = jdof_dim_cpu.numpy() if hasattr(jdof_dim_cpu, "numpy") else np.asarray(jdof_dim_cpu, dtype=int) + jt = self._to_numpy(self.model.joint_type, dtype=int) + jdof_dim = self._to_numpy(self.model.joint_dof_dim, dtype=int) dim_np = np.zeros((n_j,), dtype=np.int32) for j in range(n_j): @@ -639,11 +789,11 @@ def _init_joint_constraint_layout(self) -> None: elif jt[j] == JointType.FIXED: dim_np[j] = 2 elif jt[j] == JointType.REVOLUTE: - dim_np[j] = 3 # [linear, angular, ang_drive_limit] + dim_np[j] = 3 elif jt[j] == JointType.PRISMATIC: - dim_np[j] = 3 # [linear, angular, lin_drive_limit] + dim_np[j] = 3 elif jt[j] == JointType.D6: - dim_np[j] = 2 + int(jdof_dim[j, 0]) + int(jdof_dim[j, 1]) # [linear, angular, per-DOF drive/limit] + dim_np[j] = 2 + int(jdof_dim[j, 0]) + int(jdof_dim[j, 1]) else: if jt[j] != JointType.FREE: raise NotImplementedError( @@ -662,32 +812,18 @@ def _init_joint_constraint_layout(self) -> None: self.joint_constraint_dim = wp.array(dim_np, dtype=wp.int32, device=self.device) self.joint_constraint_start = wp.array(start_np, dtype=wp.int32, device=self.device) - def _init_joint_penalty_k(self, k_start_joint_linear: float, k_start_joint_angular: float): - """ - Build initial joint penalty state on CPU and upload to solver device. - - This initializes the solver-owned joint constraint parameter arrays used by VBD. - The arrays are sized by ``self.joint_constraint_count`` and indexed using - ``self.joint_constraint_start`` (solver constraint indexing), not by model DOF indexing. - - Arrays: - - ``k0``: initial penalty stiffness for each solver constraint scalar (stored as ``self.joint_penalty_k``) - - ``k_min``: warmstart floor for each solver constraint scalar (stored as ``self.joint_penalty_k_min``) - - ``k_max``: stiffness cap for each solver constraint scalar (stored as ``self.joint_penalty_k_max``) - - ``kd``: damping coefficient for each solver constraint scalar (stored as ``self.joint_penalty_kd``) - - Supported rigid joint constraint types in SolverVBD: - - ``JointType.CABLE`` (2 scalars: stretch + bend) - - ``JointType.BALL`` (1 scalar: isotropic linear anchor-coincidence) - - ``JointType.FIXED`` (2 scalars: isotropic linear + isotropic angular) - - ``JointType.REVOLUTE`` (3 scalars: isotropic linear + 2-DOF perpendicular angular + angular drive/limit) - - ``JointType.PRISMATIC`` (3 scalars: 2-DOF perpendicular linear + isotropic angular + linear drive/limit) - - ``JointType.D6`` (2 + lin_count + ang_count scalars: projected linear + projected angular + per-DOF drive/limit) - - Drive/limit slots use AVBD with per-mode clamping in the primal (``wp.min(avbd_ke, model_ke)``). - Drive and limit share one slot per free DOF (mutually exclusive at runtime). - - ``JointType.FREE`` joints (created by :meth:`ModelBuilder.add_body`) are not constraints and are ignored. + def _init_joint_penalty_k(self): + """Build initial joint penalty state on CPU and upload to solver device. + + Returns: + (joint_penalty_k, joint_penalty_k_min, joint_penalty_k_max) tuple: + - joint_penalty_k: mutable current penalty per constraint scalar + - joint_penalty_k_min: frozen floor (= k_start when beta > 0, slot ceiling when beta = 0) + - joint_penalty_k_max: frozen ceiling (= slot-specific ke) + + Side effects (stored on self): + - joint_penalty_kd: damping coefficient per constraint scalar + - joint_is_hard: hard/soft flag per constraint scalar (1 = hard, 0 = soft) """ if ( not hasattr(self, "joint_constraint_start") @@ -705,55 +841,40 @@ def _init_joint_penalty_k(self, k_start_joint_linear: float, k_start_joint_angul ) constraint_count = self.joint_constraint_count + lin_k_start = self.rigid_joint_linear_k_start + ang_k_start = self.rigid_joint_angular_k_start + with wp.ScopedDevice("cpu"): - # Per-constraint AVBD penalty state: - # - k0: initial penalty stiffness for this scalar constraint - # - k_min: warmstart floor (so k doesn't decay below this across steps) - # - k_max: stiffness cap (so k never exceeds the chosen target for this constraint) - # - # We start from solver-level seeds (k_start_*), but clamp to the per-constraint cap (k_max) so we always - # satisfy k_min <= k0 <= k_max. - stretch_k = max(0.0, k_start_joint_linear) - bend_k = max(0.0, k_start_joint_angular) - joint_k_min_np = np.zeros((constraint_count,), dtype=float) - joint_k0_np = np.zeros((constraint_count,), dtype=float) - # Per-constraint stiffness caps used for AVBD warmstart clamping and penalty growth limiting. - # - Cable constraints: use model.joint_target_ke (cable material/constraint tuning; still model-DOF indexed) - # - Rigid constraints (BALL/FIXED/REVOLUTE/PRISMATIC/D6): use solver-level caps (rigid_joint_linear_ke/angular_ke) - # Start from zeros and explicitly fill per joint/constraint-slot below for clarity. joint_k_max_np = np.zeros((constraint_count,), dtype=float) + joint_k_init_np = np.zeros((constraint_count,), dtype=float) joint_kd_np = np.zeros((constraint_count,), dtype=float) + is_hard_np = np.zeros((constraint_count,), dtype=np.int32) + + jt = self._to_numpy(self.model.joint_type, dtype=int) + jdofs = self._to_numpy(self.model.joint_qd_start, dtype=int) + jtarget_ke = self._to_numpy(self.model.joint_target_ke, dtype=float) + jtarget_kd = self._to_numpy(self.model.joint_target_kd, dtype=float) + jlimit_ke = self._to_numpy(self.model.joint_limit_ke, dtype=float) + jdof_dim = self._to_numpy(self.model.joint_dof_dim, dtype=int) + jc_start = self._to_numpy(self.joint_constraint_start, dtype=np.int32) + + # Per-joint hard/soft mode from model attribute (default=1, hard). + vbd_attrs: Any = getattr(self.model, "vbd", None) + if vbd_attrs is not None and hasattr(vbd_attrs, "joint_is_hard"): + j_is_hard = self._to_numpy(vbd_attrs.joint_is_hard, dtype=np.int32) + if not np.all((j_is_hard == 0) | (j_is_hard == 1)): + raise ValueError("model.vbd.joint_is_hard values must be 0 (soft) or 1 (hard).") + else: + j_is_hard = np.ones(self.model.joint_count, dtype=np.int32) - jt_cpu = self.model.joint_type.to("cpu") - jdofs_cpu = self.model.joint_qd_start.to("cpu") - jtarget_ke_cpu = self.model.joint_target_ke.to("cpu") - jtarget_kd_cpu = self.model.joint_target_kd.to("cpu") - jlimit_ke_cpu = self.model.joint_limit_ke.to("cpu") - jdof_dim_cpu = self.model.joint_dof_dim.to("cpu") - jc_start_cpu = self.joint_constraint_start.to("cpu") - - jt = jt_cpu.numpy() if hasattr(jt_cpu, "numpy") else np.asarray(jt_cpu, dtype=int) - jdofs = jdofs_cpu.numpy() if hasattr(jdofs_cpu, "numpy") else np.asarray(jdofs_cpu, dtype=int) - jc_start = ( - jc_start_cpu.numpy() if hasattr(jc_start_cpu, "numpy") else np.asarray(jc_start_cpu, dtype=np.int32) - ) - jtarget_ke = ( - jtarget_ke_cpu.numpy() if hasattr(jtarget_ke_cpu, "numpy") else np.asarray(jtarget_ke_cpu, dtype=float) - ) - jtarget_kd = ( - jtarget_kd_cpu.numpy() if hasattr(jtarget_kd_cpu, "numpy") else np.asarray(jtarget_kd_cpu, dtype=float) - ) - jlimit_ke = ( - jlimit_ke_cpu.numpy() if hasattr(jlimit_ke_cpu, "numpy") else np.asarray(jlimit_ke_cpu, dtype=float) - ) - jdof_dim = jdof_dim_cpu.numpy() if hasattr(jdof_dim_cpu, "numpy") else np.asarray(jdof_dim_cpu, dtype=int) + structural_linear_ke = self.rigid_joint_linear_ke + structural_angular_ke = self.rigid_joint_angular_ke n_j = self.model.joint_count for j in range(n_j): if jt[j] == JointType.CABLE: c0 = int(jc_start[j]) dof0 = int(jdofs[j]) - # CABLE requires 2 DOF entries in model.joint_target_ke/kd starting at joint_qd_start[j]. if dof0 < 0 or (dof0 + 1) >= len(jtarget_ke) or (dof0 + 1) >= len(jtarget_kd): raise RuntimeError( "SolverVBD _init_joint_penalty_k: JointType.CABLE requires 2 DOF entries in " @@ -761,146 +882,118 @@ def _init_joint_penalty_k(self, k_start_joint_linear: float, k_start_joint_angul f"Got joint_index={j}, joint_qd_start={dof0}, " f"len(joint_target_ke)={len(jtarget_ke)}, len(joint_target_kd)={len(jtarget_kd)}." ) - # Constraint 0: cable stretch; constraint 1: cable bend - # Caps come from model.joint_target_ke (still model DOF indexed for cable material tuning). - joint_k_max_np[c0] = jtarget_ke[dof0] - joint_k_max_np[c0 + 1] = jtarget_ke[dof0 + 1] - # Per-slot warmstart lower bounds: - # - Use k_start_* as the floor, but clamp to the cap so k_min <= k_max always. - joint_k_min_np[c0] = min(stretch_k, joint_k_max_np[c0]) - joint_k_min_np[c0 + 1] = min(bend_k, joint_k_max_np[c0 + 1]) - # Initial seed: clamp to cap so k0 <= k_max - joint_k0_np[c0] = min(stretch_k, joint_k_max_np[c0]) - joint_k0_np[c0 + 1] = min(bend_k, joint_k_max_np[c0 + 1]) - # Damping comes from model.joint_target_kd (still model DOF indexed for cable tuning). + ke_stretch = jtarget_ke[dof0] + ke_bend = jtarget_ke[dof0 + 1] + joint_k_max_np[c0] = ke_stretch + joint_k_max_np[c0 + 1] = ke_bend + joint_k_init_np[c0] = ke_stretch if lin_k_start is None else min(lin_k_start, ke_stretch) + joint_k_init_np[c0 + 1] = ke_bend if ang_k_start is None else min(ang_k_start, ke_bend) joint_kd_np[c0] = jtarget_kd[dof0] joint_kd_np[c0 + 1] = jtarget_kd[dof0 + 1] elif jt[j] == JointType.BALL: - # BALL joints: isotropic linear anchor-coincidence constraint stored as a single scalar. c0 = int(jc_start[j]) - joint_k_max_np[c0] = self.rigid_joint_linear_ke - k_floor = min(stretch_k, self.rigid_joint_linear_ke) - joint_k_min_np[c0] = k_floor - joint_k0_np[c0] = k_floor + joint_k_max_np[c0] = structural_linear_ke + joint_k_init_np[c0] = ( + structural_linear_ke if lin_k_start is None else min(lin_k_start, structural_linear_ke) + ) joint_kd_np[c0] = self.rigid_joint_linear_kd + is_hard_np[c0] = int(j_is_hard[j]) elif jt[j] == JointType.FIXED: - # FIXED joints are enforced as: - # - 1 isotropic linear anchor-coincidence constraint (vector error, scalar penalty) - # - 1 isotropic angular constraint (rotation-vector error, scalar penalty) c0 = int(jc_start[j]) - - # Linear cap + floor (isotropic) - joint_k_max_np[c0 + 0] = self.rigid_joint_linear_ke - k_lin_floor = min(stretch_k, self.rigid_joint_linear_ke) - joint_k_min_np[c0 + 0] = k_lin_floor - joint_k0_np[c0 + 0] = k_lin_floor + joint_k_max_np[c0 + 0] = structural_linear_ke + joint_k_init_np[c0 + 0] = ( + structural_linear_ke if lin_k_start is None else min(lin_k_start, structural_linear_ke) + ) joint_kd_np[c0 + 0] = self.rigid_joint_linear_kd - - # Angular cap + floor (isotropic) - joint_k_max_np[c0 + 1] = self.rigid_joint_angular_ke - k_ang_floor = min(bend_k, self.rigid_joint_angular_ke) - joint_k_min_np[c0 + 1] = k_ang_floor - joint_k0_np[c0 + 1] = k_ang_floor + is_hard_np[c0 + 0] = int(j_is_hard[j]) + joint_k_max_np[c0 + 1] = structural_angular_ke + joint_k_init_np[c0 + 1] = ( + structural_angular_ke if ang_k_start is None else min(ang_k_start, structural_angular_ke) + ) joint_kd_np[c0 + 1] = self.rigid_joint_angular_kd + is_hard_np[c0 + 1] = int(j_is_hard[j]) elif jt[j] == JointType.REVOLUTE: - # REVOLUTE joints: isotropic linear + 2-DOF perpendicular angular + angular drive/limit c0 = int(jc_start[j]) - - joint_k_max_np[c0 + 0] = self.rigid_joint_linear_ke - k_lin_floor = min(stretch_k, self.rigid_joint_linear_ke) - joint_k_min_np[c0 + 0] = k_lin_floor - joint_k0_np[c0 + 0] = k_lin_floor + joint_k_max_np[c0 + 0] = structural_linear_ke + joint_k_init_np[c0 + 0] = ( + structural_linear_ke if lin_k_start is None else min(lin_k_start, structural_linear_ke) + ) joint_kd_np[c0 + 0] = self.rigid_joint_linear_kd - - joint_k_max_np[c0 + 1] = self.rigid_joint_angular_ke - k_ang_floor = min(bend_k, self.rigid_joint_angular_ke) - joint_k_min_np[c0 + 1] = k_ang_floor - joint_k0_np[c0 + 1] = k_ang_floor + is_hard_np[c0 + 0] = int(j_is_hard[j]) + joint_k_max_np[c0 + 1] = structural_angular_ke + joint_k_init_np[c0 + 1] = ( + structural_angular_ke if ang_k_start is None else min(ang_k_start, structural_angular_ke) + ) joint_kd_np[c0 + 1] = self.rigid_joint_angular_kd - - # Drive/limit slot for free angular DOF (slot c0 + 2). - # Drive and limit share one AVBD slot (mutually exclusive at runtime). - # Per-mode clamping in the primal prevents branch-switch overshoot. + is_hard_np[c0 + 1] = int(j_is_hard[j]) dof0 = int(jdofs[j]) dl_k_max = max(float(jtarget_ke[dof0]), float(jlimit_ke[dof0])) - dl_seed = min(bend_k, dl_k_max) # angular DOF -> bend_k seed + dl_seed = dl_k_max if ang_k_start is None else min(ang_k_start, dl_k_max) joint_k_max_np[c0 + 2] = dl_k_max - joint_k_min_np[c0 + 2] = dl_seed - joint_k0_np[c0 + 2] = dl_seed - joint_kd_np[c0 + 2] = 0.0 # damping is non-adaptive, read from model in primal + joint_k_init_np[c0 + 2] = dl_seed + joint_kd_np[c0 + 2] = 0.0 elif jt[j] == JointType.PRISMATIC: - # PRISMATIC joints: 2-DOF perpendicular linear + isotropic angular + linear drive/limit c0 = int(jc_start[j]) - - joint_k_max_np[c0 + 0] = self.rigid_joint_linear_ke - k_lin_floor = min(stretch_k, self.rigid_joint_linear_ke) - joint_k_min_np[c0 + 0] = k_lin_floor - joint_k0_np[c0 + 0] = k_lin_floor + joint_k_max_np[c0 + 0] = structural_linear_ke + joint_k_init_np[c0 + 0] = ( + structural_linear_ke if lin_k_start is None else min(lin_k_start, structural_linear_ke) + ) joint_kd_np[c0 + 0] = self.rigid_joint_linear_kd - - joint_k_max_np[c0 + 1] = self.rigid_joint_angular_ke - k_ang_floor = min(bend_k, self.rigid_joint_angular_ke) - joint_k_min_np[c0 + 1] = k_ang_floor - joint_k0_np[c0 + 1] = k_ang_floor + is_hard_np[c0 + 0] = int(j_is_hard[j]) + joint_k_max_np[c0 + 1] = structural_angular_ke + joint_k_init_np[c0 + 1] = ( + structural_angular_ke if ang_k_start is None else min(ang_k_start, structural_angular_ke) + ) joint_kd_np[c0 + 1] = self.rigid_joint_angular_kd - - # Drive/limit slot for free linear DOF (slot c0 + 2). + is_hard_np[c0 + 1] = int(j_is_hard[j]) dof0 = int(jdofs[j]) dl_k_max = max(float(jtarget_ke[dof0]), float(jlimit_ke[dof0])) - dl_seed = min(stretch_k, dl_k_max) # linear DOF -> stretch_k seed + dl_seed = dl_k_max if lin_k_start is None else min(lin_k_start, dl_k_max) joint_k_max_np[c0 + 2] = dl_k_max - joint_k_min_np[c0 + 2] = dl_seed - joint_k0_np[c0 + 2] = dl_seed + joint_k_init_np[c0 + 2] = dl_seed joint_kd_np[c0 + 2] = 0.0 elif jt[j] == JointType.D6: - # D6 joints: projected linear + projected angular + per-DOF drive/limit c0 = int(jc_start[j]) dof0 = int(jdofs[j]) - lc = int(jdof_dim[j, 0]) # free linear DOF count - ac = int(jdof_dim[j, 1]) # free angular DOF count - - joint_k_max_np[c0 + 0] = self.rigid_joint_linear_ke - k_lin_floor = min(stretch_k, self.rigid_joint_linear_ke) - joint_k_min_np[c0 + 0] = k_lin_floor - joint_k0_np[c0 + 0] = k_lin_floor + lc = int(jdof_dim[j, 0]) + ac = int(jdof_dim[j, 1]) + joint_k_max_np[c0 + 0] = structural_linear_ke + joint_k_init_np[c0 + 0] = ( + structural_linear_ke if lin_k_start is None else min(lin_k_start, structural_linear_ke) + ) joint_kd_np[c0 + 0] = self.rigid_joint_linear_kd - - joint_k_max_np[c0 + 1] = self.rigid_joint_angular_ke - k_ang_floor = min(bend_k, self.rigid_joint_angular_ke) - joint_k_min_np[c0 + 1] = k_ang_floor - joint_k0_np[c0 + 1] = k_ang_floor + is_hard_np[c0 + 0] = int(j_is_hard[j]) + joint_k_max_np[c0 + 1] = structural_angular_ke + joint_k_init_np[c0 + 1] = ( + structural_angular_ke if ang_k_start is None else min(ang_k_start, structural_angular_ke) + ) joint_kd_np[c0 + 1] = self.rigid_joint_angular_kd - - # Per free linear DOF drive/limit slots + is_hard_np[c0 + 1] = int(j_is_hard[j]) for li in range(lc): dof_idx = dof0 + li slot = c0 + 2 + li dl_k_max = max(float(jtarget_ke[dof_idx]), float(jlimit_ke[dof_idx])) - dl_seed = min(stretch_k, dl_k_max) + dl_seed = dl_k_max if lin_k_start is None else min(lin_k_start, dl_k_max) joint_k_max_np[slot] = dl_k_max - joint_k_min_np[slot] = dl_seed - joint_k0_np[slot] = dl_seed + joint_k_init_np[slot] = dl_seed joint_kd_np[slot] = 0.0 - - # Per free angular DOF drive/limit slots for ai in range(ac): dof_idx = dof0 + lc + ai slot = c0 + 2 + lc + ai dl_k_max = max(float(jtarget_ke[dof_idx]), float(jlimit_ke[dof_idx])) - dl_seed = min(bend_k, dl_k_max) + dl_seed = dl_k_max if ang_k_start is None else min(ang_k_start, dl_k_max) joint_k_max_np[slot] = dl_k_max - joint_k_min_np[slot] = dl_seed - joint_k0_np[slot] = dl_seed + joint_k_init_np[slot] = dl_seed joint_kd_np[slot] = 0.0 else: - # Layout builder already validated supported types; nothing to do for FREE. pass - # Upload to device: initial penalties, per-constraint caps, damping, and warmstart floors. - self.joint_penalty_k_min = wp.array(joint_k_min_np, dtype=float, device=self.device) - self.joint_penalty_k_max = wp.array(joint_k_max_np, dtype=float, device=self.device) self.joint_penalty_kd = wp.array(joint_kd_np, dtype=float, device=self.device) - return wp.array(joint_k0_np, dtype=float, device=self.device) + self.joint_is_hard = wp.array(is_hard_np, dtype=wp.int32, device=self.device) + k = wp.array(joint_k_init_np, dtype=float, device=self.device) + k_min = wp.array(joint_k_init_np.copy(), dtype=float, device=self.device) + k_max = wp.array(joint_k_max_np, dtype=float, device=self.device) + return k, k_min, k_max def _init_joint_rest_angle(self): """Compute per-DOF rest-pose joint angles from ``model.joint_q``. @@ -915,25 +1008,17 @@ def _init_joint_rest_angle(self): Only angular DOFs of REVOLUTE and D6 joints need nonzero entries. Linear DOFs (PRISMATIC, D6 linear) use absolute geometric measurements (``d_along``) and - are unaffected — their entries are left at 0. + are unaffected - their entries are left at 0. """ dof_count = self.model.joint_dof_count rest_angle_np = np.zeros(dof_count, dtype=float) with wp.ScopedDevice("cpu"): - jt_cpu = self.model.joint_type.to("cpu") - jq_cpu = self.model.joint_q.to("cpu") - jq_start_cpu = self.model.joint_q_start.to("cpu") - jqd_start_cpu = self.model.joint_qd_start.to("cpu") - jdof_dim_cpu = self.model.joint_dof_dim.to("cpu") - - jt = jt_cpu.numpy() if hasattr(jt_cpu, "numpy") else np.asarray(jt_cpu, dtype=int) - jq = jq_cpu.numpy() if hasattr(jq_cpu, "numpy") else np.asarray(jq_cpu, dtype=float) - jq_start = jq_start_cpu.numpy() if hasattr(jq_start_cpu, "numpy") else np.asarray(jq_start_cpu, dtype=int) - jqd_start = ( - jqd_start_cpu.numpy() if hasattr(jqd_start_cpu, "numpy") else np.asarray(jqd_start_cpu, dtype=int) - ) - jdof_dim = jdof_dim_cpu.numpy() if hasattr(jdof_dim_cpu, "numpy") else np.asarray(jdof_dim_cpu, dtype=int) + jt = self._to_numpy(self.model.joint_type, dtype=int) + jq = self._to_numpy(self.model.joint_q, dtype=float) + jq_start = self._to_numpy(self.model.joint_q_start, dtype=int) + jqd_start = self._to_numpy(self.model.joint_qd_start, dtype=int) + jdof_dim = self._to_numpy(self.model.joint_dof_dim, dtype=int) for j in range(self.model.joint_count): if jt[j] == JointType.REVOLUTE: @@ -950,60 +1035,14 @@ def _init_joint_rest_angle(self): return wp.array(rest_angle_np, dtype=float, device=self.device) - def _init_dahl_params(self, eps_max_input, tau_input, model): - """ - Initialize per-joint Dahl friction parameters. - - Args: - eps_max_input: float or array-like. Maximum strain (curvature) [rad]. - - Scalar: broadcast to all joints - - Array-like (length = model.joint_count): per-joint values - - Per-joint disable: set value to 0 for that joint - tau_input: float or array-like. Memory decay length [rad]. - - Scalar: broadcast to all joints - - Array-like (length = model.joint_count): per-joint values - - Per-joint disable: set value to 0 for that joint - model: Model object - - Notes: - - This function validates shapes and converts to device arrays; it does not clamp or validate ranges. - Kernels perform any necessary early-outs based on zero values. - - To disable Dahl friction: - - Globally: pass enable_dahl_friction=False to the constructor - - Per-joint: set dahl_eps_max=0 or dahl_tau=0 for those joints - """ - n = model.joint_count - - # eps_max - if isinstance(eps_max_input, (int, float)): - self.joint_dahl_eps_max = wp.full(n, eps_max_input, dtype=float, device=self.device) - else: - # Convert to numpy first - x = eps_max_input.to("cpu") if hasattr(eps_max_input, "to") else eps_max_input - eps_np = x.numpy() if hasattr(x, "numpy") else np.asarray(x, dtype=float) - if eps_np.shape[0] != n: - raise ValueError(f"dahl_eps_max length {eps_np.shape[0]} != joint_count {n}") - # Direct host-to-device copy - self.joint_dahl_eps_max = wp.array(eps_np, dtype=float, device=self.device) - - # tau - if isinstance(tau_input, (int, float)): - self.joint_dahl_tau = wp.full(n, tau_input, dtype=float, device=self.device) - else: - # Convert to numpy first - x = tau_input.to("cpu") if hasattr(tau_input, "to") else tau_input - tau_np = x.numpy() if hasattr(x, "numpy") else np.asarray(x, dtype=float) - if tau_np.shape[0] != n: - raise ValueError(f"dahl_tau length {tau_np.shape[0]} != joint_count {n}") - # Direct host-to-device copy - self.joint_dahl_tau = wp.array(tau_np, dtype=float, device=self.device) - @override @classmethod def register_custom_attributes(cls, builder: ModelBuilder) -> None: """Register solver-specific custom Model attributes for SolverVBD. - Currently used for cable bending plasticity/hysteresis (Dahl friction model). + Currently used for: + - Cable bending plasticity/hysteresis (Dahl friction model) + - Per-joint structural constraint mode (hard/soft) Attributes are declared in the ``vbd`` namespace so they can be authored in scenes and in USD as ``newton:vbd:``. @@ -1028,6 +1067,16 @@ def register_custom_attributes(cls, builder: ModelBuilder) -> None: namespace="vbd", ) ) + builder.add_custom_attribute( + ModelBuilder.CustomAttribute( + name="joint_is_hard", + frequency=Model.AttributeFrequency.JOINT, + assignment=Model.AttributeAssignment.MODEL, + dtype=wp.int32, + default=1, + namespace="vbd", + ) + ) # ===================================================== # Adjacency Building Methods @@ -1323,16 +1372,102 @@ def _compute_rigid_force_element_adjacency(self, model): # ===================================================== def set_rigid_history_update(self, update: bool): - """Set whether the next step() should update rigid solver history (warmstarts). + """Set whether the next step() should update rigid solver history. + + When True (default), the step refreshes contact state from the Contacts + buffer: restores state from the hash table, initializes penalty_k/lambda/C0, + and rebuilds per-body contact lists. When False, the step reuses the + current contact state and runs a light snapshot to keep the hash table fresh. - This setting applies only to the next call to :meth:`step` and is then reset to ``True``. - This is useful for substepping, where history update frequency might differ from the - simulation step frequency (e.g. updating only on the first substep). + Joint AVBD maintenance (C0 snapshot, lambda decay) + runs every step regardless of this flag via step_joint_C0_lambda(). + + This setting applies only to the next call to :meth:`step` and is then + reset to True. Useful for substepping where collision detection frequency + differs from the simulation step frequency. Args: - update: If True, update rigid warmstart state. If False, reuse previous. + update: If True, update rigid solver state. If False, reuse previous. """ - self.update_rigid_history = update + self._update_rigid_history = update + + def set_joint_constraint_mode(self, joint_index: int, hard: bool, slot: int | None = None): + """Set hard or soft constraint mode for a joint's structural slots at runtime. + + Hard mode (augmented Lagrangian): uses persistent lambda + C0 stabilization + to drive constraint violation toward zero across iterations. + Soft mode (penalty-only): uses penalty stiffness only (no lambda or C0 state). + + Structural slots are LINEAR (slot 0) and ANGULAR (slot 1). Drive/limit slots + (slot 2+) are always soft and cannot be set to hard. + + For bulk initialization at build time (avoids per-joint roundtrips), use + the ``joint_is_hard`` model custom attribute instead:: + + SolverVBD.register_custom_attributes(builder) # before adding joints + ... + model = builder.finalize() + model.vbd.joint_is_hard.numpy()[j] = 0 # set joint j to soft + solver = SolverVBD(model, ...) + + Args: + joint_index: Index of the joint to modify. + hard: True for hard mode (AL), False for soft mode (penalty-only). + slot: Specific slot index to set. If None, sets all structural slots. + Use JointSlot.LINEAR / JointSlot.ANGULAR (equivalently + JointSlot.STRETCH / JointSlot.BEND for cables). + + Raises: + ValueError: If the joint index is out of range, or the slot is a + drive/limit slot (>= 2), or the slot exceeds the joint's + constraint dimension. + """ + n_j = self.model.joint_count + if joint_index < 0 or joint_index >= n_j: + raise ValueError(f"joint_index={joint_index} out of range [0, {n_j}).") + + with wp.ScopedDevice("cpu"): + c_start_np = self._to_numpy(self.joint_constraint_start, dtype=np.int32) + c_dim_np = self._to_numpy(self.joint_constraint_dim, dtype=np.int32) + is_hard_np = self._to_numpy(self.joint_is_hard, dtype=np.int32) + + c0 = int(c_start_np[joint_index]) + cdim = int(c_dim_np[joint_index]) + val = 1 if hard else 0 + + if slot is not None: + if slot < 0 or slot >= 2: + raise ValueError( + f"Cannot set hard mode on slot={slot}. " + "Only structural slots (LINEAR=0, ANGULAR=1) support hard mode." + ) + if slot >= cdim: + raise ValueError( + f"slot={slot} exceeds joint constraint dimension ({cdim}) for joint_index={joint_index}." + ) + is_hard_np[c0 + slot] = val + else: + structural_count = min(cdim, 2) + for s in range(structural_count): + is_hard_np[c0 + s] = val + + self.joint_is_hard = wp.array(is_hard_np, dtype=wp.int32, device=self.device) + + if not hard: + lam_lin_np = self._to_numpy(self.joint_lambda_lin) + lam_ang_np = self._to_numpy(self.joint_lambda_ang) + C0_lin_np = self._to_numpy(self.joint_C0_lin) + C0_ang_np = self._to_numpy(self.joint_C0_ang) + if slot is None or slot == 0: + lam_lin_np[joint_index] = [0.0, 0.0, 0.0] + C0_lin_np[joint_index] = [0.0, 0.0, 0.0] + if (slot is None or slot == 1) and cdim > 1: + lam_ang_np[joint_index] = [0.0, 0.0, 0.0] + C0_ang_np[joint_index] = [0.0, 0.0, 0.0] + self.joint_lambda_lin = wp.array(lam_lin_np, dtype=wp.vec3, device=self.device) + self.joint_lambda_ang = wp.array(lam_ang_np, dtype=wp.vec3, device=self.device) + self.joint_C0_lin = wp.array(C0_lin_np, dtype=wp.vec3, device=self.device) + self.joint_C0_ang = wp.array(C0_ang_np, dtype=wp.vec3, device=self.device) @override def step( @@ -1346,13 +1481,14 @@ def step( """Execute one simulation timestep using VBD (particles) and AVBD (rigid bodies). The solver follows a 3-phase structure: - 1. Initialize: Forward integrate particles and rigid bodies, detect collisions, warmstart penalties + 1. Initialize: Forward integrate particles and rigid bodies, detect collisions, initialize contact state 2. Iterate: Interleave particle VBD iterations and rigid body AVBD iterations 3. Finalize: Update velocities and persistent state (Dahl friction) - To control rigid substepping behavior (warmstart history), call - :meth:`set_rigid_history_update` - before calling this method. It defaults to ``True`` and is reset to ``True`` after each call. + To control rigid body substepping behavior, call set_rigid_history_update(). + When True (default), the step rebuilds rigid contact adjacency, re-initializes + rigid contact state (penalty_k, lambda, C0), and restores from history if enabled. + When False, reuses previous rigid contact state. Resets to True after each call. Args: state_in: Input state. @@ -1363,23 +1499,106 @@ def step( depend on this argument. dt: Time step size. """ - # Use and reset the rigid history update flag (warmstarts). - update_rigid_history = self.update_rigid_history - self.update_rigid_history = True + update_rigid = self._update_rigid_history + self._update_rigid_history = True if control is None: control = self.model.control(clone_variables=False) - self._initialize_rigid_bodies(state_in, control, contacts, dt, update_rigid_history) + update_rigid = self._initialize_rigid_bodies(state_in, control, contacts, dt, update_rigid) self._initialize_particles(state_in, state_out, dt) for iter_num in range(self.iterations): - self._solve_rigid_body_iteration(state_in, state_out, control, contacts, dt) + self._solve_rigid_body_iteration(state_in, control, contacts, dt) self._solve_particle_iteration(state_in, state_out, contacts, dt, iter_num) - self._finalize_rigid_bodies(state_out, dt) + # Snapshot rigid contact history before rigid finalize advances body_q_prev. + self._snapshot_rigid_contact_history(contacts, update_rigid) + self._finalize_rigid_bodies( + state_in, state_out, dt, apply_stick_deadzone=contacts is not None and self.rigid_contact_hard + ) self._finalize_particles(state_out, dt) + def _snapshot_rigid_contact_history(self, contacts: Contacts | None, refresh: bool): + """Write contact solver state to the hash table. + + When contact history is enabled: + - On refresh steps: full snapshot into the hash table and cache slot indices. + - On non-refresh steps: light update of lambda/stick/penalty_k at cached slot indices. + In soft mode, lambda and stick_flag are inert but stored uniformly. + """ + if not self.rigid_contact_history or contacts is None: + return + + if self.model.body_count == 0 or self.integrate_with_external_rigid_solver: + return + + contact_launch_dim = contacts.rigid_contact_max + if refresh: + ht = self.contact_history_ht + if self._rigid_contact_match_max_age == 0: + ht.clear() + else: + wp.launch( + kernel=evict_contact_history, + dim=ht.capacity, + inputs=[ + ht.keys, + self.contact_history_age, + self._rigid_contact_match_max_age, + ], + device=self.device, + ) + ht.active_slots.zero_() + wp.launch( + kernel=snapshot_contact_history, + dim=contact_launch_dim, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + contacts.rigid_contact_normal, + self.body_body_contact_lambda, + self.body_body_contact_stick_flag, + self.body_body_contact_penalty_k, + ht.keys, + ht.active_slots, + self.rigid_contact_match_cell_size_inv, + ], + outputs=[ + self.contact_history_point0, + self.contact_history_point1, + self.contact_history_lambda, + self.contact_history_normal, + self.contact_history_stick_flag, + self.contact_history_penalty_k, + self.contact_history_age, + self.contact_history_slots, + ], + device=self.device, + ) + + else: + wp.launch( + kernel=snapshot_contact_history_light, + dim=contact_launch_dim, + inputs=[ + contacts.rigid_contact_count, + self.contact_history_slots, + self.body_body_contact_lambda, + self.body_body_contact_stick_flag, + self.body_body_contact_penalty_k, + ], + outputs=[ + self.contact_history_lambda, + self.contact_history_stick_flag, + self.contact_history_penalty_k, + ], + device=self.device, + ) + def _penetration_free_truncation(self, particle_q_out=None): """ Modify displacements_in in-place, also modify particle_q if its not None @@ -1457,8 +1676,6 @@ def _initialize_particles(self, state_in: State, state_out: State, dt: float): self.pos_prev_collision_detection.assign(state_in.particle_q) self.particle_displacements.zero_() - model = self.model - wp.launch( kernel=forward_step, inputs=[ @@ -1487,17 +1704,21 @@ def _initialize_rigid_bodies( control: Control, contacts: Contacts | None, dt: float, - update_rigid_history: bool, - ): + refresh: bool, + ) -> bool: """Initialize rigid body states for AVBD solver (pre-iteration phase). Performs forward integration and initializes contact-related AVBD state when contacts are provided. - If ``contacts`` is None, rigid contact-related work is skipped: - no per-body contact adjacency is built, and no contact penalties are warmstarted. + If contacts is None, rigid contact-related work is skipped: + no per-body contact adjacency is built, and no contact state is initialized or restored. - If ``control`` provides ``joint_f``, per-DOF joint forces are mapped to body spatial + If control provides joint_f, per-DOF joint forces are mapped to body spatial wrenches and included in the forward integration (shifting the inertial target). + + Returns: + The effective refresh flag (may be promoted from False to True when + contact state needs first-time allocation or resizing). """ model = self.model @@ -1505,11 +1726,158 @@ def _initialize_rigid_bodies( # Rigid-only initialization # --------------------------- if model.body_count > 0 and not self.integrate_with_external_rigid_solver: - # Accumulate per-DOF joint forces (joint_f) into body spatial wrenches. - # Clone body_f to avoid mutating user state; the clone is used only for integration. + # Force refresh when contact state is not yet allocated or undersized. + if ( + not refresh + and contacts is not None + and contacts.rigid_contact_max > 0 + and self.body_body_contact_penalty_k.shape[0] < contacts.rigid_contact_max + ): + refresh = True + + # Contact C0 + history restore BEFORE integration: body_q is the collide frame + # for all bodies (dynamic and kinematic) at this point. + if refresh: + if contacts is None: + self.body_body_contact_counts.zero_() + else: + contact_launch_dim = contacts.rigid_contact_max + + if self.body_body_contact_penalty_k.shape[0] < contact_launch_dim: + self._init_body_body_contact_state(contact_launch_dim) + + # Build body-body contact lists + self.body_body_contact_counts.zero_() + self.body_body_contact_overflow_max.zero_() + wp.launch( + kernel=build_body_body_contact_lists, + dim=contact_launch_dim, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + model.shape_body, + self.body_body_contact_buffer_pre_alloc, + ], + outputs=[ + self.body_body_contact_counts, + self.body_body_contact_indices, + self.body_body_contact_overflow_max, + ], + device=self.device, + ) + wp.launch( + kernel=check_contact_overflow, + dim=1, + inputs=[self.body_body_contact_overflow_max, self.body_body_contact_buffer_pre_alloc, 0], + device=self.device, + ) + + # Restore AVBD body-body contact state from history and pre-compute material properties + if self.rigid_contact_history: + if ( + self.contact_history_ht is None + or self.contact_history_slots.shape[0] < contacts.rigid_contact_max + ): + self._init_contact_history(contacts.rigid_contact_max) + + wp.launch( + kernel=init_body_body_contacts_avbd, + dim=contact_launch_dim, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + contacts.rigid_contact_normal, + model.shape_material_ke, + model.shape_material_kd, + model.shape_material_mu, + self.rigid_contact_hard, + self.contact_history_ht.keys, + self.contact_history_point0, + self.contact_history_point1, + self.contact_history_lambda, + self.contact_history_normal, + self.contact_history_stick_flag, + self.contact_history_penalty_k, + self.rigid_contact_match_cell_size_inv, + self.rigid_contact_match_tolerance_sq, + self.rigid_contact_k_start_value, + ], + outputs=[ + self.body_body_contact_penalty_k, + self.body_body_contact_lambda, + self.body_body_contact_material_kd, + self.body_body_contact_material_mu, + self.body_body_contact_material_ke, + ], + device=self.device, + ) + else: + wp.launch( + kernel=init_body_body_contact_materials, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + model.shape_material_ke, + model.shape_material_kd, + model.shape_material_mu, + self.rigid_contact_k_start_value, + ], + outputs=[ + self.body_body_contact_penalty_k, + self.body_body_contact_material_kd, + self.body_body_contact_material_mu, + self.body_body_contact_material_ke, + ], + dim=contact_launch_dim, + device=self.device, + ) + self.body_body_contact_lambda.zero_() + + # Per-step k decay + lambda decay + C0 (body_q is still collide frame here). + if contacts is not None and contacts.rigid_contact_max > 0: + contact_launch_dim = contacts.rigid_contact_max + gamma_lambda = ( + self.rigid_avbd_gamma if (self.rigid_contact_history and self.rigid_contact_hard) else 0.0 + ) + wp.launch( + kernel=step_body_body_contact_C0_lambda, + dim=contact_launch_dim, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + contacts.rigid_contact_normal, + contacts.rigid_contact_margin0, + contacts.rigid_contact_margin1, + model.shape_body, + state_in.body_q, + self.rigid_contact_hard, + gamma_lambda, + self.rigid_avbd_gamma, + self.body_body_contact_material_ke, + self.rigid_contact_k_start_value, + ], + outputs=[ + self.body_body_contact_penalty_k, + self.body_body_contact_C0, + self.body_body_contact_lambda, + ], + device=self.device, + ) + self.body_body_contact_stick_flag.zero_() + + # Accumulate joint_f into body wrenches (scratch buffer avoids mutating user state). body_f_for_integration = state_in.body_f if model.joint_count > 0 and control is not None and control.joint_f is not None: - body_f_for_integration = wp.clone(state_in.body_f) + wp.copy(self._body_f_for_integration, state_in.body_f) + body_f_for_integration = self._body_f_for_integration wp.launch( kernel=apply_joint_forces, dim=model.joint_count, @@ -1533,7 +1901,7 @@ def _initialize_rigid_bodies( device=self.device, ) - # Forward integrate rigid bodies (snapshots body_q_prev for dynamic bodies only) + # Forward integrate rigid bodies (body_q modified in-place for dynamic bodies only). wp.launch( kernel=forward_step_rigid_bodies, inputs=[ @@ -1550,74 +1918,39 @@ def _initialize_rigid_bodies( ], outputs=[ self.body_inertia_q, - self.body_q_prev, ], dim=model.body_count, device=self.device, ) - if update_rigid_history: - # Contact warmstarts / adjacency are optional: skip completely if contacts=None. - if contacts is not None: - # Use the Contacts buffer capacity as launch dimension - contact_launch_dim = contacts.rigid_contact_max - - # Build per-body contact lists once per step - # Build body-body (rigid-rigid) contact lists - self.body_body_contact_counts.zero_() - wp.launch( - kernel=build_body_body_contact_lists, - dim=contact_launch_dim, - inputs=[ - contacts.rigid_contact_count, - contacts.rigid_contact_shape0, - contacts.rigid_contact_shape1, - model.shape_body, - self.body_body_contact_buffer_pre_alloc, - ], - outputs=[ - self.body_body_contact_counts, - self.body_body_contact_indices, - ], - device=self.device, - ) - - # Warmstart AVBD body-body contact penalties and pre-compute material properties - wp.launch( - kernel=warmstart_body_body_contacts, - inputs=[ - contacts.rigid_contact_count, - contacts.rigid_contact_shape0, - contacts.rigid_contact_shape1, - model.shape_material_ke, - model.shape_material_kd, - model.shape_material_mu, - self.k_start_body_contact, - ], - outputs=[ - self.body_body_contact_penalty_k, - self.body_body_contact_material_ke, - self.body_body_contact_material_kd, - self.body_body_contact_material_mu, - ], - dim=contact_launch_dim, - device=self.device, - ) - - # Warmstart AVBD penalty parameters for joints using the same cadence - # as rigid history updates. - if model.joint_count > 0: - wp.launch( - kernel=warmstart_joints, - inputs=[ - self.joint_penalty_k_max, - self.joint_penalty_k_min, - self.avbd_gamma, - self.joint_penalty_k, # input/output - ], - dim=self.joint_constraint_count, - device=self.device, - ) + if model.joint_count > 0: + wp.launch( + kernel=step_joint_C0_lambda, + dim=model.joint_count, + inputs=[ + model.joint_enabled, + model.joint_parent, + model.joint_child, + model.joint_X_p, + model.joint_X_c, + self.body_q_prev, + model.body_q, + self.joint_constraint_start, + self.joint_constraint_dim, + self.joint_is_hard, + self.rigid_avbd_gamma, + ], + outputs=[ + self.joint_penalty_k, + self.joint_penalty_k_min, + self.joint_penalty_k_max, + self.joint_C0_lin, + self.joint_C0_ang, + self.joint_lambda_lin, + self.joint_lambda_ang, + ], + device=self.device, + ) # Compute Dahl hysteresis parameters for cable bending (once per timestep, frozen during iterations) if self.enable_dahl_friction and model.joint_count > 0: @@ -1631,9 +1964,9 @@ def _initialize_rigid_bodies( model.joint_X_p, model.joint_X_c, self.joint_constraint_start, - self.joint_penalty_k_max, - self.body_q_prev, # Use previous body transforms (start of step) for linearization - model.body_q, # rest body transforms + self.joint_penalty_k, + self.body_q_prev, + model.body_q, self.joint_sigma_prev, self.joint_kappa_prev, self.joint_dkappa_prev, @@ -1651,13 +1984,21 @@ def _initialize_rigid_bodies( # --------------------------- # Body-particle interaction # --------------------------- - if model.particle_count > 0 and update_rigid_history and contacts is not None: - # Build body-particle (rigid-particle) contact lists only when SolverVBD - # is integrating rigid bodies itself; the external rigid solver path - # does not use these per-body adjacency structures. Also skip if there - # are no rigid bodies in the model. + particle_refresh = refresh + if ( + not particle_refresh + and model.particle_count > 0 + and contacts is not None + and contacts.soft_contact_max > 0 + and self.body_particle_contact_penalty_k.shape[0] < contacts.soft_contact_max + ): + particle_refresh = True + + if model.particle_count > 0 and particle_refresh and contacts is not None: + # Build body-particle contact lists (only when SolverVBD integrates bodies). if not self.integrate_with_external_rigid_solver and model.body_count > 0: self.body_particle_contact_counts.zero_() + self.body_particle_contact_overflow_max.zero_() wp.launch( kernel=build_body_particle_contact_lists, dim=contacts.soft_contact_max, @@ -1670,16 +2011,23 @@ def _initialize_rigid_bodies( outputs=[ self.body_particle_contact_counts, self.body_particle_contact_indices, + self.body_particle_contact_overflow_max, ], device=self.device, ) + wp.launch( + kernel=check_contact_overflow, + dim=1, + inputs=[self.body_particle_contact_overflow_max, self.body_particle_contact_buffer_pre_alloc, 1], + device=self.device, + ) - # Warmstart AVBD body-particle contact penalties and pre-compute material properties. - # This is useful both when SolverVBD integrates rigid bodies and when an external - # rigid solver is used, since cloth-rigid soft contacts still rely on these penalties. + # Init body-particle material properties (needed for both internal and external rigid solver). soft_contact_launch_dim = contacts.soft_contact_max + if self.body_particle_contact_penalty_k.shape[0] < soft_contact_launch_dim: + self._init_body_particle_contact_state(soft_contact_launch_dim) wp.launch( - kernel=warmstart_body_particle_contacts, + kernel=init_body_particle_contacts, inputs=[ contacts.soft_contact_count, contacts.soft_contact_shape, @@ -1689,18 +2037,20 @@ def _initialize_rigid_bodies( model.shape_material_ke, model.shape_material_kd, model.shape_material_mu, - self.k_start_body_contact, + self.rigid_contact_k_start_value, ], outputs=[ self.body_particle_contact_penalty_k, - self.body_particle_contact_material_ke, self.body_particle_contact_material_kd, self.body_particle_contact_material_mu, + self.body_particle_contact_material_ke, ], dim=soft_contact_launch_dim, device=self.device, ) + return refresh + def _solve_particle_iteration( self, state_in: State, state_out: State, contacts: Contacts | None, dt: float, iter_num: int ): @@ -1889,9 +2239,7 @@ def _solve_particle_iteration( wp.copy(state_out.particle_q, state_in.particle_q) - def _solve_rigid_body_iteration( - self, state_in: State, state_out: State, control: Control, contacts: Contacts | None, dt: float - ): + def _solve_rigid_body_iteration(self, state_in: State, control: Control, contacts: Contacts | None, dt: float): """Solve one AVBD iteration for rigid bodies (per-iteration phase). Accumulates contact and joint forces/hessians, solves 6x6 rigid body systems per color, @@ -1899,44 +2247,8 @@ def _solve_rigid_body_iteration( """ model = self.model - # Early-return path: - # - If rigid bodies are integrated by an external solver, skip the AVBD rigid-body solve but still - # update body-particle soft-contact penalties so adaptive stiffness is used for particle-shape - # interaction. - # - If there are no rigid bodies at all (body_count == 0), we likewise skip the rigid-body solve, - # but must still update particle-shape soft-contact penalties (e.g., particles colliding with the - # ground plane where shape_body == -1). skip_rigid_solve = self.integrate_with_external_rigid_solver or model.body_count == 0 if skip_rigid_solve: - if model.particle_count > 0 and contacts is not None: - # Use external rigid poses when enabled; otherwise use the current VBD poses. - body_q = state_out.body_q if self.integrate_with_external_rigid_solver else state_in.body_q - - # Model.state() leaves State.body_q as None when body_count == 0. Warp kernels still - # require a wp.array argument; for static shapes (shape_body == -1) the kernel never - # indexes this array, so an empty placeholder is sufficient. - if body_q is None: - body_q = self._empty_body_q - - wp.launch( - kernel=update_duals_body_particle_contacts, - dim=contacts.soft_contact_max, - inputs=[ - contacts.soft_contact_count, - contacts.soft_contact_particle, - contacts.soft_contact_shape, - contacts.soft_contact_body_pos, - contacts.soft_contact_normal, - state_in.particle_q, - model.particle_radius, - model.shape_body, - body_q, - self.body_particle_contact_material_ke, - self.avbd_beta, - self.body_particle_contact_penalty_k, # input/output - ], - device=self.device, - ) return # Zero out forces and hessians @@ -1952,8 +2264,7 @@ def _solve_rigid_body_iteration( for color in range(len(body_color_groups)): color_group = body_color_groups[color] - # Gauss-Seidel contact accumulation: evaluate contacts for bodies in this color - # Accumulate body-particle forces and Hessians on bodies (per-body, per-color) + # Accumulate body-particle contact forces/hessians for bodies in this color if model.particle_count > 0 and contacts is not None: wp.launch( kernel=accumulate_body_particle_contacts_per_body, @@ -1961,32 +2272,22 @@ def _solve_rigid_body_iteration( inputs=[ dt, color_group, - # particle state state_in.particle_q, self.particle_q_prev, model.particle_radius, - # rigid body state self.body_q_prev, state_in.body_q, - state_in.body_qd, model.body_com, self.body_inv_mass_effective, - # AVBD body-particle soft contact penalties and material properties self.friction_epsilon, self.body_particle_contact_penalty_k, self.body_particle_contact_material_kd, self.body_particle_contact_material_mu, - # soft contact data (body-particle contacts) contacts.soft_contact_count, contacts.soft_contact_particle, - contacts.soft_contact_shape, contacts.soft_contact_body_pos, contacts.soft_contact_body_vel, contacts.soft_contact_normal, - # shape/material data - model.shape_material_mu, - model.shape_body, - # per-body adjacency (body-particle contacts) self.body_particle_contact_buffer_pre_alloc, self.body_particle_contact_counts, self.body_particle_contact_indices, @@ -2017,6 +2318,10 @@ def _solve_rigid_body_iteration( self.body_body_contact_penalty_k, self.body_body_contact_material_kd, self.body_body_contact_material_mu, + self.body_body_contact_lambda, + self.body_body_contact_C0, + self.rigid_contact_alpha, + self.rigid_contact_hard, contacts.rigid_contact_count, contacts.rigid_contact_shape0, contacts.rigid_contact_shape1, @@ -2067,16 +2372,20 @@ def _solve_rigid_body_iteration( self.joint_penalty_kd, self.joint_sigma_start, self.joint_C_fric, - # Drive parameters (DOF-indexed) model.joint_target_ke, model.joint_target_kd, control.joint_target_pos, control.joint_target_vel, - # Limit parameters (DOF-indexed) model.joint_limit_lower, model.joint_limit_upper, model.joint_limit_ke, model.joint_limit_kd, + self.joint_lambda_lin, + self.joint_lambda_ang, + self.joint_C0_lin, + self.joint_C0_ang, + self.joint_is_hard, + self.rigid_joint_alpha, model.joint_dof_dim, self.joint_rest_angle, self.body_forces, @@ -2086,23 +2395,13 @@ def _solve_rigid_body_iteration( self.body_hessian_aa, ], outputs=[ - state_out.body_q, + state_in.body_q, ], dim=color_group.size, device=self.device, ) - wp.launch( - kernel=copy_rigid_body_transforms_back, - inputs=[color_group, state_out.body_q], - outputs=[state_in.body_q], - dim=color_group.size, - device=self.device, - ) - if contacts is not None: - # AVBD dual update: update adaptive penalties based on constraint violation - # Update body-body (rigid-rigid) contact penalties contact_launch_dim = contacts.rigid_contact_max wp.launch( kernel=update_duals_body_body_contacts, @@ -2117,16 +2416,25 @@ def _solve_rigid_body_iteration( contacts.rigid_contact_margin0, contacts.rigid_contact_margin1, model.shape_body, - state_out.body_q, + state_in.body_q, + self.body_q_prev, + self.body_body_contact_material_mu, + self.body_body_contact_C0, + self.rigid_contact_alpha, + self.rigid_contact_hard, + self.body_inv_mass_effective, self.body_body_contact_material_ke, - self.avbd_beta, - self.body_body_contact_penalty_k, # input/output + self.rigid_linear_beta, + self.body_body_contact_penalty_k, + self.body_body_contact_lambda, + ], + outputs=[ + self.body_body_contact_stick_flag, ], device=self.device, ) - # Update body-particle contact penalties - if model.particle_count > 0: + if model.particle_count > 0 and not self.integrate_with_external_rigid_solver: soft_contact_launch_dim = contacts.soft_contact_max wp.launch( kernel=update_duals_body_particle_contacts, @@ -2140,17 +2448,16 @@ def _solve_rigid_body_iteration( state_in.particle_q, model.particle_radius, model.shape_body, - # Rigid poses come from SolverVBD itself when - # integrate_with_external_rigid_solver=False state_in.body_q, self.body_particle_contact_material_ke, - self.avbd_beta, - self.body_particle_contact_penalty_k, # input/output + self.rigid_linear_beta, + ], + outputs=[ + self.body_particle_contact_penalty_k, ], device=self.device, ) - # Update joint penalties at new positions if model.joint_count > 0: wp.launch( kernel=update_duals_joint, @@ -2165,33 +2472,42 @@ def _solve_rigid_body_iteration( model.joint_axis, model.joint_qd_start, self.joint_constraint_start, - self.joint_penalty_k_max, - state_out.body_q, + state_in.body_q, model.body_q, - self.avbd_beta, - self.joint_penalty_k, # input/output model.joint_dof_dim, - self.joint_rest_angle, - # Drive/limit parameters for adaptive drive/limit penalty growth + self.joint_C0_lin, + self.joint_C0_ang, + self.joint_is_hard, + self.rigid_joint_alpha, + self.joint_penalty_k_max, + self.rigid_linear_beta, + self.rigid_angular_beta, model.joint_target_ke, control.joint_target_pos, model.joint_limit_lower, model.joint_limit_upper, model.joint_limit_ke, + self.joint_rest_angle, + self.joint_penalty_k, + self.joint_lambda_lin, + self.joint_lambda_ang, ], device=self.device, ) def collect_rigid_contact_forces( - self, state: State, contacts: Contacts | None, dt: float + self, body_q: wp.array, body_q_prev: wp.array, contacts: Contacts | None, dt: float ) -> tuple[wp.array, wp.array, wp.array, wp.array, wp.array, wp.array]: """Collect per-contact rigid contact forces and world-space application points. - This produces a **contact-specific** buffer that coupling code can filter (e.g., proxy contacts only). - Args: - state (State): Simulation state containing rigid body transforms/velocities - used for contact-force evaluation. + body_q (wp.array[wp.transform]): Current body transforms (world frame), + typically ``state_out.body_q`` after a ``step()`` call. + body_q_prev (wp.array[wp.transform]): Previous body transforms (world frame) + corresponding to the start of the step. The caller must snapshot + ``solver.body_q_prev`` **before** calling ``step()``, because + ``step()`` advances the solver's internal ``body_q_prev`` to the + end-of-step pose. contacts (Optional[Contacts]): Contact data buffers containing rigid contact geometry/material references. If None, the function returns default zero/sentinel outputs. @@ -2213,8 +2529,8 @@ def collect_rigid_contact_forces( - force_on_body1: Contact force applied to body1 in world frame, wp.vec3 [N]. - rigid_contact_count: Length-1 active rigid-contact count, int32. """ - # Allocate/resize persistent buffers to match contact capacity. max_contacts = int(contacts.rigid_contact_shape0.shape[0]) if contacts is not None else 0 + if not hasattr(self, "_rigid_contact_body0") or self._rigid_contact_body0 is None: self._rigid_contact_body0 = None @@ -2224,10 +2540,13 @@ def collect_rigid_contact_forces( self._rigid_contact_point0_world = wp.zeros(max_contacts, dtype=wp.vec3, device=self.device) self._rigid_contact_point1_world = wp.zeros(max_contacts, dtype=wp.vec3, device=self.device) + # Ensure buffers are large enough so the kernel can run even before the first step(). + if max_contacts > 0 and self.body_body_contact_penalty_k.shape[0] < max_contacts: + self._init_body_body_contact_state(max_contacts) + missing_rigid_state = any( - arr is None + arr is None or arr.shape[0] == 0 for arr in ( - getattr(self, "body_q_prev", None), getattr(self, "body_body_contact_penalty_k", None), getattr(self, "body_body_contact_material_kd", None), getattr(self, "body_body_contact_material_mu", None), @@ -2239,18 +2558,12 @@ def collect_rigid_contact_forces( contacts.rigid_contact_force.zero_() if no_active_contacts or missing_rigid_state: - # Keep outputs in a known default state for coupling paths where rigid AVBD - # internal buffers are not initialized (e.g., external rigid solver mode). self._rigid_contact_body0 = wp.full(max_contacts, -1, dtype=wp.int32, device=self.device) self._rigid_contact_body1 = wp.full(max_contacts, -1, dtype=wp.int32, device=self.device) self._rigid_contact_point0_world = wp.zeros(max_contacts, dtype=wp.vec3, device=self.device) self._rigid_contact_point1_world = wp.zeros(max_contacts, dtype=wp.vec3, device=self.device) - rigid_contact_count = ( - contacts.rigid_contact_count - if contacts is not None and contacts.rigid_contact_count is not None - else wp.zeros(1, dtype=wp.int32, device=self.device) - ) + rigid_contact_count = wp.zeros(1, dtype=wp.int32, device=self.device) return ( self._rigid_contact_body0, self._rigid_contact_body1, @@ -2265,8 +2578,6 @@ def collect_rigid_contact_forces( # Type narrowing: remaining path requires a valid Contacts instance. assert contacts is not None - # Reuse the existing per-contact force buffer in Contacts (allocated by default). - # Force convention: force is applied to body1, and -force is applied to body0. wp.launch( kernel=compute_rigid_contact_forces, dim=max_contacts, @@ -2281,12 +2592,16 @@ def collect_rigid_contact_forces( contacts.rigid_contact_margin0, contacts.rigid_contact_margin1, self.model.shape_body, - state.body_q, - self.body_q_prev, + body_q, + body_q_prev, self.model.body_com, self.body_body_contact_penalty_k, self.body_body_contact_material_kd, self.body_body_contact_material_mu, + self.body_body_contact_lambda, + self.body_body_contact_C0, + self.rigid_contact_alpha, + self.rigid_contact_hard, float(self.friction_epsilon), ], outputs=[ @@ -2321,10 +2636,13 @@ def _finalize_particles(self, state_out: State, dt: float): device=self.device, ) - def _finalize_rigid_bodies(self, state_out: State, dt: float): + def _finalize_rigid_bodies(self, state_in: State, state_out: State, dt: float, apply_stick_deadzone: bool): """Finalize rigid body velocities and Dahl friction state after AVBD iterations (post-iteration phase). Updates rigid body velocities using BDF1 and updates Dahl hysteresis state for cable bending. + Also transfers the final body poses from state_in to state_out. When requested, + the fused finalize kernel first applies the body-level stick-contact deadzone + before computing velocity from the accepted pose. """ model = self.model @@ -2332,20 +2650,25 @@ def _finalize_rigid_bodies(self, state_out: State, dt: float): if model.body_count == 0 or self.integrate_with_external_rigid_solver: return - # Velocity update (BDF1) after all iterations wp.launch( kernel=update_body_velocity, inputs=[ dt, - state_out.body_q, + state_in.body_q, model.body_com, + self.body_body_contact_buffer_pre_alloc, + self.body_body_contact_counts, + self.body_body_contact_indices, + self.body_body_contact_stick_flag, + int(apply_stick_deadzone), + self.rigid_contact_stick_freeze_translation_eps, + self.rigid_contact_stick_freeze_angular_eps, ], - outputs=[self.body_q_prev, state_out.body_qd], + outputs=[self.body_q_prev, state_out.body_qd, state_in.body_qd, state_out.body_q], dim=model.body_count, device=self.device, ) - # Update Dahl hysteresis state after solver convergence (for next timestep's memory) if self.enable_dahl_friction and model.joint_count > 0: wp.launch( kernel=update_cable_dahl_state, @@ -2357,14 +2680,15 @@ def _finalize_rigid_bodies(self, state_out: State, dt: float): model.joint_X_p, model.joint_X_c, self.joint_constraint_start, - self.joint_penalty_k_max, + self.joint_penalty_k, + self.joint_is_hard, state_out.body_q, model.body_q, self.joint_dahl_eps_max, self.joint_dahl_tau, - self.joint_sigma_prev, # input/output - self.joint_kappa_prev, # input/output - self.joint_dkappa_prev, # input/output + self.joint_sigma_prev, + self.joint_kappa_prev, + self.joint_dkappa_prev, ], dim=model.joint_count, device=self.device, diff --git a/newton/examples/basic/example_basic_conveyor.py b/newton/examples/basic/example_basic_conveyor.py index b807f50146..5137f993c9 100644 --- a/newton/examples/basic/example_basic_conveyor.py +++ b/newton/examples/basic/example_basic_conveyor.py @@ -183,20 +183,20 @@ def __init__(self, viewer, args=None): belt_cfg = newton.ModelBuilder.ShapeConfig( mu=1.2, - ke=1.0e7, # vbd only - kd=1.0e-5, # vbd only + ke=1.0e4, # vbd only + kd=0.0, # vbd only collision_group=BELT_COLLISION_GROUP, ) rail_cfg = newton.ModelBuilder.ShapeConfig( mu=0.8, - ke=1.0e7, # vbd only - kd=1.0e-5, # vbd only + ke=1.0e4, # vbd only + kd=0.0, # vbd only collision_group=RAIL_COLLISION_GROUP, ) bag_cfg = newton.ModelBuilder.ShapeConfig( mu=1.0, - ke=1.0e7, # vbd only - kd=1.0e-5, # vbd only + ke=1.0e4, # vbd only + kd=0.0, # vbd only restitution=0.0, ) @@ -327,7 +327,7 @@ def __init__(self, viewer, args=None): solver_type = getattr(args, "solver", "xpbd") if args is not None else "xpbd" if solver_type == "vbd": - self.solver = newton.solvers.SolverVBD(self.model) + self.solver = newton.solvers.SolverVBD(self.model, rigid_body_contact_buffer_size=512) else: self.solver = newton.solvers.SolverXPBD(self.model) diff --git a/newton/examples/basic/example_basic_urdf.py b/newton/examples/basic/example_basic_urdf.py index c87c80117b..6253c11318 100644 --- a/newton/examples/basic/example_basic_urdf.py +++ b/newton/examples/basic/example_basic_urdf.py @@ -41,11 +41,10 @@ def __init__(self, viewer, args): quadruped.default_joint_cfg.armature = 0.01 if self.solver_type == "vbd": - quadruped.default_joint_cfg.target_ke = 2000.0 - quadruped.default_joint_cfg.target_kd = 1.0e-3 - quadruped.default_joint_cfg.limit_kd = 1.0e-5 - quadruped.default_shape_cfg.ke = 1.0e7 - quadruped.default_shape_cfg.kd = 1.0e-1 + quadruped.default_joint_cfg.target_ke = 1.0e4 + quadruped.default_joint_cfg.target_kd = 0.0 + quadruped.default_shape_cfg.ke = 1.0e5 + quadruped.default_shape_cfg.kd = 0.0 quadruped.default_shape_cfg.mu = 1.0 else: quadruped.default_joint_cfg.target_ke = 2000.0 @@ -87,11 +86,16 @@ def __init__(self, viewer, args): newton.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, self.model) if self.solver_type == "vbd": - self.update_step_interval = 10 - self.solver = newton.solvers.SolverVBD(self.model, iterations=1, rigid_contact_k_start=1.0e6) + self.update_step_interval = 1 + self.solver = newton.solvers.SolverVBD( + self.model, + iterations=1, + # Example-specific stabilization for low-iteration AVBD; prefer more iterations / non-zero damping. + rigid_avbd_gamma=0.9, + ) else: self.update_step_interval = 1 - self.solver = newton.solvers.SolverXPBD(self.model) + self.solver = newton.solvers.SolverXPBD(self.model, iterations=1) self.state_0 = self.model.state() self.state_1 = self.model.state() @@ -119,12 +123,12 @@ def simulate(self): # apply forces to the model self.viewer.apply_forces(self.state_0) - update_step_history = (substep % self.update_step_interval) == 0 - if update_step_history: + refresh_contacts = (substep % self.update_step_interval) == 0 + if refresh_contacts: self.model.collide(self.state_0, self.contacts) if self.solver_type == "vbd": - self.solver.set_rigid_history_update(update_step_history) + self.solver.set_rigid_history_update(refresh_contacts) self.solver.step(self.state_0, self.state_1, self.control, self.contacts, self.sim_dt) diff --git a/newton/examples/cable/example_cable_bundle_hysteresis.py b/newton/examples/cable/example_cable_bundle_hysteresis.py index cc5e958a5d..f4af98696c 100644 --- a/newton/examples/cable/example_cable_bundle_hysteresis.py +++ b/newton/examples/cable/example_cable_bundle_hysteresis.py @@ -162,7 +162,7 @@ def __init__( self.frame_dt = 1.0 / self.fps self.sim_time = 0.0 self.sim_substeps = 10 - self.sim_iterations = 10 + self.sim_iterations = 5 self.update_step_interval = 10 self.sim_dt = self.frame_dt / self.sim_substeps @@ -174,20 +174,20 @@ def __init__( self.cable_gap_multiplier = 1.1 bend_stiffness = 1.0e1 bend_damping = 5.0e-2 - stretch_stiffness = 1.0e6 - stretch_damping = 0.0 builder = newton.ModelBuilder() builder.rigid_gap = 0.05 - # Register solver-specific custom attributes (Dahl plasticity parameters live on the Model) - newton.solvers.SolverVBD.register_custom_attributes(builder) + # Register solver-specific custom attributes (Dahl plasticity parameters live on the Model). + # SolverVBD auto-detects these and enables Dahl friction when present. + if with_dahl: + newton.solvers.SolverVBD.register_custom_attributes(builder) builder.gravity = -9.81 # Set default material properties for cables (cable-to-cable contact) - builder.default_shape_cfg.ke = 1.0e6 # Contact stiffness - builder.default_shape_cfg.kd = 1.0e-2 # Contact damping - builder.default_shape_cfg.mu = 2.0 # Friction coefficient + builder.default_shape_cfg.ke = 1.0e5 # Contact stiffness + builder.default_shape_cfg.kd = 0.0 + builder.default_shape_cfg.mu = 1.0e0 # Friction coefficient # Bundle layout: align cable center with obstacle center # Obstacles span x in [0.5, 2.5], center at x=1.5 @@ -219,16 +219,12 @@ def __init__( radius=self.cable_radius, bend_stiffness=bend_stiffness, bend_damping=bend_damping, - stretch_stiffness=stretch_stiffness, - stretch_damping=stretch_damping, label=f"bundle_cable_{i}", ) # Create moving obstacles (capsules arranged along X axis) obstacle_cfg = newton.ModelBuilder.ShapeConfig( density=builder.default_shape_cfg.density, - ke=1.0e6, - kd=1.0e-2, kf=builder.default_shape_cfg.kf, ka=builder.default_shape_cfg.ka, mu=0.0, # Frictionless obstacles @@ -270,8 +266,6 @@ def __init__( # Add ground plane ground_cfg = newton.ModelBuilder.ShapeConfig( density=builder.default_shape_cfg.density, - ke=1.0e4, - kd=1.0e-1, kf=builder.default_shape_cfg.kf, ka=builder.default_shape_cfg.ka, mu=2.5, @@ -286,17 +280,14 @@ def __init__( self.model = builder.finalize() # Author Dahl friction parameters (per-joint) via custom model attributes. - # These are read by SolverVBD when rigid_enable_dahl_friction=True. - if hasattr(self.model, "vbd"): + # SolverVBD auto-detects these and enables Dahl friction when present. + if with_dahl and hasattr(self.model, "vbd"): self.model.vbd.dahl_eps_max.fill_(float(eps_max)) self.model.vbd.dahl_tau.fill_(float(tau)) - # Create VBD solver with Dahl friction (cable bending hysteresis) self.solver = newton.solvers.SolverVBD( self.model, iterations=self.sim_iterations, - friction_epsilon=0.1, - rigid_enable_dahl_friction=with_dahl, ) # Initialize states and contacts @@ -378,15 +369,12 @@ def simulate(self): device=self.solver.device, ) - # Decide whether to refresh rigid solver history (anchors used for long-range joint damping) - # and recompute contacts on this substep, using a configurable cadence. - update_step_history = (substep % self.update_step_interval) == 0 - - # Collide for contact detection - if update_step_history: + # Collision detection and contact refresh cadence. + refresh_contacts = (substep % self.update_step_interval) == 0 + if refresh_contacts: self.model.collide(self.state_0, self.contacts) - self.solver.set_rigid_history_update(update_step_history) + self.solver.set_rigid_history_update(refresh_contacts) self.solver.step( self.state_0, self.state_1, diff --git a/newton/examples/cable/example_cable_pile.py b/newton/examples/cable/example_cable_pile.py index 4f085b3489..97dcae8db1 100644 --- a/newton/examples/cable/example_cable_pile.py +++ b/newton/examples/cable/example_cable_pile.py @@ -29,7 +29,6 @@ def __init__( slope_angle_deg: float = 20.0, slope_mu: float | None = None, ): - # Store viewer and arguments self.viewer = viewer self.args = args @@ -43,27 +42,25 @@ def __init__( # Cable pile parameters self.num_elements = 40 - segment_length = 0.05 # 5cm segments + segment_length = 0.05 self.cable_length = self.num_elements * segment_length cable_radius = 0.012 # Layers and lanes - layers = 4 + layers = 10 lanes_per_layer = 10 - # Increase spacing to accommodate lateral waviness without initial intersections lane_spacing = max(8.0 * cable_radius, 0.15) layer_gap = cable_radius * 6.0 builder = newton.ModelBuilder() - builder.rigid_gap = 0.05 # Default for all shapes + builder.rigid_gap = 0.0 rod_bodies_all: list[int] = [] - # Set default material properties before adding any shapes - # Default config will be used by plane and any shape without explicit cfg - builder.default_shape_cfg.ke = 1.0e6 # Contact stiffness (used by plane) - builder.default_shape_cfg.kd = 1.0e-1 # Contact damping - builder.default_shape_cfg.mu = 5.0 # Friction coefficient + # Material properties + builder.default_shape_cfg.mu = 1.0e0 + builder.default_shape_cfg.ke = 1.0e5 + builder.default_shape_cfg.kd = 1.0e-4 cable_shape_cfg = newton.ModelBuilder.ShapeConfig( density=builder.default_shape_cfg.density, @@ -75,12 +72,11 @@ def __init__( restitution=builder.default_shape_cfg.restitution, ) - # Generate a ground plane (optionally sloped for friction tests) + # Ground plane (optionally sloped for friction tests) if slope_enabled: angle = math.radians(slope_angle_deg) rot = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), angle) - # Optionally override friction for the slope without changing defaults for other shapes slope_cfg = builder.default_shape_cfg if slope_mu is not None: slope_cfg = newton.ModelBuilder.ShapeConfig( @@ -101,13 +97,17 @@ def __init__( cfg=slope_cfg, ) else: - builder.add_ground_plane() + ground_cfg = newton.ModelBuilder.ShapeConfig( + mu=1.0e9, + ke=builder.default_shape_cfg.ke, + kd=builder.default_shape_cfg.kd, + ) + builder.add_ground_plane(cfg=ground_cfg) - # Build layered lanes of cables with alternating orientations (no initial intersections) + # Build layered lanes of cables with alternating orientations for layer in range(layers): orient = "x" if (layer % 2 == 0) else "y" z0 = 0.3 + layer * layer_gap - # Generate parallel lanes along the orthogonal axis only for lane in range(lanes_per_layer): offset = (lane - (lanes_per_layer - 1) * 0.5) * lane_spacing if orient == "x": @@ -115,14 +115,12 @@ def __init__( else: start = wp.vec3(offset, 0.0, z0) - # Regular waviness and no twist for repeatable layout across layers wav = 0.5 twist = 0.0 dir_vec = wp.vec3(1.0, 0.0, 0.0) if orient == "x" else wp.vec3(0.0, 1.0, 0.0) ortho_vec = wp.vec3(0.0, 1.0, 0.0) if orient == "x" else wp.vec3(1.0, 0.0, 0.0) - # Center the cable around start to avoid initial overlaps with neighbors. cable_length = float(self.cable_length) start0 = start - 0.5 * cable_length * dir_vec pts = newton.utils.create_straight_cable_points( @@ -132,9 +130,9 @@ def __init__( num_segments=int(self.num_elements), ) - # Apply sinusoidal waviness along an orthogonal axis (optional). + # Sinusoidal waviness along orthogonal axis cycles = 2.0 - waviness_scale = 0.05 # amplitude fraction of length when wav=1.0 + waviness_scale = 0.05 if wav > 0.0: for i in range(len(pts)): t = i / self.num_elements @@ -149,36 +147,36 @@ def __init__( quaternions=edge_q, radius=cable_radius, cfg=cable_shape_cfg, - bend_stiffness=1.0e1, - bend_damping=5.0e-1, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-4, + bend_stiffness=1.0e0, + bend_damping=1.0e0, label=f"cable_l{layer}_{lane}", ) rod_bodies_all.extend(rod_bodies) - # Color bodies for VBD solver builder.color() - # Finalize model self.model = builder.finalize() - # Create VBD solver for rigid body simulation self.solver = newton.solvers.SolverVBD( self.model, iterations=self.sim_iterations, - friction_epsilon=0.1, + rigid_body_contact_buffer_size=256, + rigid_contact_history=True, ) - # Initialize states and contacts self.state_0 = self.model.state() self.state_1 = self.model.state() self.control = self.model.control() - self.contacts = self.model.contacts() + self.viewer.set_model(self.model) - # Optional capture for CUDA + if hasattr(self.viewer, "picking"): + ps = self.viewer.picking.pick_state.numpy() + ps[0]["pick_stiffness"] = 20.0 + ps[0]["pick_damping"] = 0.0 + self.viewer.picking.pick_state.assign(ps) + self.capture() def capture(self): @@ -194,11 +192,7 @@ def simulate(self): """Execute all simulation substeps for one frame.""" for _substep in range(self.sim_substeps): self.state_0.clear_forces() - - # Apply forces to the model self.viewer.apply_forces(self.state_0) - - # Collide for contact detection self.model.collide(self.state_0, self.contacts) self.solver.step( @@ -209,11 +203,10 @@ def simulate(self): self.sim_dt, ) - # Swap states self.state_0, self.state_1 = self.state_1, self.state_0 def step(self): - """Advance simulation by one frame (either via CUDA graph or direct execution).""" + """Advance simulation by one frame.""" if self.graph: wp.capture_launch(self.graph) else: @@ -231,30 +224,22 @@ def render(self): def test_final(self): """Test cable pile simulation for stability and correctness (called after simulation).""" cable_radius = 0.012 - cable_diameter = 2.0 * cable_radius # 0.024m - layers = 4 + cable_diameter = 2.0 * cable_radius + layers = 10 - # Use same tolerance for ground penetration and pile height offset - tolerance = 0.1 # Soft contacts allow some penetration and gaps + tolerance = 0.5 - # Calculate maximum expected height for SETTLED pile - # After gravity settles the pile, cables should be stacked: - # 4 layers = 4 cable diameters high (approximately) - # Plus compression tolerance and contact gaps max_z_settled = layers * cable_diameter + tolerance ground_tolerance = tolerance - # Check final state after viewer has run 100 frames (no additional simulation needed) if self.state_0.body_q is not None and self.state_0.body_qd is not None: body_positions = self.state_0.body_q.numpy() body_velocities = self.state_0.body_qd.numpy() - # Test 1: Check for numerical stability assert np.isfinite(body_positions).all(), "Non-finite positions" assert np.isfinite(body_velocities).all(), "Non-finite velocities" - # Test 2: Check physical bounds - cables should stay within pile height - z_positions = body_positions[:, 2] # Z positions (Newton uses Z-up) + z_positions = body_positions[:, 2] min_z = np.min(z_positions) max_z_actual = np.max(z_positions) @@ -263,19 +248,13 @@ def test_final(self): ) assert max_z_actual < max_z_settled, ( f"Pile too high: max_z={max_z_actual:.3f} > expected {max_z_settled:.3f} " - f"(4 layers x {cable_diameter:.3f}m diameter + tolerance)" + f"({layers} layers x {cable_diameter:.3f}m diameter + tolerance)" ) - # Test 3: Velocity should be reasonable (pile shouldn't explode) assert (np.abs(body_velocities) < 5e2).all(), "Velocities too large" if __name__ == "__main__": - # Parse arguments and initialize viewer viewer, args = newton.examples.init() - - # Create example and run - # For sloped friction tests, use: Example(viewer, args, slope_enabled=True, slope_angle_deg=20.0, slope_mu=0.8) example = Example(viewer, args) - newton.examples.run(example, args) diff --git a/newton/examples/cable/example_cable_twist.py b/newton/examples/cable/example_cable_twist.py index 0590159244..fa2452a3bd 100644 --- a/newton/examples/cable/example_cable_twist.py +++ b/newton/examples/cable/example_cable_twist.py @@ -116,8 +116,8 @@ def __init__(self, viewer, args): self.frame_dt = 1.0 / self.fps self.sim_time = 0.0 self.sim_substeps = 10 - self.sim_iterations = 2 - self.update_step_interval = 5 + self.sim_iterations = 5 + self.update_step_interval = 10 self.sim_dt = self.frame_dt / self.sim_substeps # Cable parameters @@ -126,9 +126,10 @@ def __init__(self, viewer, args): self.cable_length = self.num_elements * segment_length cable_radius = 0.02 + stretch_stiffness = 1.0e5 + # Stiffness sweep (increasing) for bend stiffness bend_stiffness_values = [1.0e1, 1.0e2, 1.0e3] - stretch_stiffness = 1.0e6 # All cables start untwisted, will be spun dynamically self.num_cables = len(bend_stiffness_values) @@ -138,7 +139,7 @@ def __init__(self, viewer, args): # Set default material properties before adding any shapes builder.default_shape_cfg.ke = 1.0e4 # Contact stiffness - builder.default_shape_cfg.kd = 1.0e-1 # Contact damping + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0e0 # Friction coefficient kinematic_body_indices = [] @@ -167,10 +168,9 @@ def __init__(self, viewer, args): positions=cable_points, quaternions=cable_edge_q, radius=cable_radius, + stretch_stiffness=stretch_stiffness, bend_stiffness=bend_stiffness, bend_damping=1.0e-2, - stretch_stiffness=stretch_stiffness, - stretch_damping=1.0e-4, label=f"cable_{i}", ) @@ -198,7 +198,7 @@ def __init__(self, viewer, args): # Finalize model self.model = builder.finalize() - self.solver = newton.solvers.SolverVBD(self.model, iterations=self.sim_iterations, friction_epsilon=0.1) + self.solver = newton.solvers.SolverVBD(self.model, iterations=self.sim_iterations) self.state_0 = self.model.state() self.state_1 = self.model.state() @@ -239,15 +239,12 @@ def simulate(self): # Apply forces to the model self.viewer.apply_forces(self.state_0) - # Decide whether to refresh solver history (anchors used for long-range damping) - # and recompute contacts on this substep, using a configurable cadence. - update_step_history = (substep % self.update_step_interval) == 0 - - # Collide for contact detection - if update_step_history: + # Collision detection and contact refresh cadence. + refresh_contacts = (substep % self.update_step_interval) == 0 + if refresh_contacts: self.model.collide(self.state_0, self.contacts) - self.solver.set_rigid_history_update(update_step_history) + self.solver.set_rigid_history_update(refresh_contacts) self.solver.step( self.state_0, self.state_1, diff --git a/newton/examples/cable/example_cable_y_junction.py b/newton/examples/cable/example_cable_y_junction.py index 51ac343e59..a5a1ae81be 100644 --- a/newton/examples/cable/example_cable_y_junction.py +++ b/newton/examples/cable/example_cable_y_junction.py @@ -38,7 +38,7 @@ def __init__(self, viewer, args): self.fps = 60 self.frame_dt = 1.0 / self.fps self.sim_time = 0.0 - self.sim_substeps = 10 + self.sim_substeps = 5 self.sim_iterations = 5 self.sim_dt = self.frame_dt / self.sim_substeps @@ -47,14 +47,12 @@ def __init__(self, viewer, args): num_segments_per_branch = 20 segment_length = 0.03 - bend_stiffness = 1.0e0 + bend_stiffness = 1.0e2 bend_damping = 1.0e-1 - stretch_stiffness = 1.0e9 - stretch_damping = 0.0 builder = newton.ModelBuilder() builder.default_shape_cfg.ke = 1.0e4 - builder.default_shape_cfg.kd = 1.0e-1 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 cable_cfg = builder.default_shape_cfg.copy() @@ -82,8 +80,6 @@ def __init__(self, viewer, args): edges=edges, radius=cable_radius, cfg=cable_cfg, - stretch_stiffness=stretch_stiffness, - stretch_damping=stretch_damping, bend_stiffness=bend_stiffness, bend_damping=bend_damping, label="y_graph", @@ -111,7 +107,6 @@ def __init__(self, viewer, args): self.solver = newton.solvers.SolverVBD( self.model, iterations=self.sim_iterations, - friction_epsilon=float(getattr(args, "friction_epsilon", 0.1)), ) self.state_0 = self.model.state() diff --git a/newton/examples/cloth/example_cloth_franka.py b/newton/examples/cloth/example_cloth_franka.py index 9772dfcce2..d50c61f824 100644 --- a/newton/examples/cloth/example_cloth_franka.py +++ b/newton/examples/cloth/example_cloth_franka.py @@ -254,7 +254,6 @@ def __init__(self, viewer, args): particle_vertex_contact_buffer_size=16, particle_edge_contact_buffer_size=20, particle_collision_detection_interval=-1, - rigid_contact_k_start=self.soft_contact_ke, ) self.viewer.set_model(self.model) diff --git a/newton/examples/contacts/example_contacts_rj45_plug.py b/newton/examples/contacts/example_contacts_rj45_plug.py index d0f3e2d946..a1252c2f8e 100644 --- a/newton/examples/contacts/example_contacts_rj45_plug.py +++ b/newton/examples/contacts/example_contacts_rj45_plug.py @@ -26,10 +26,13 @@ from newton.math import quat_between_vectors_robust from newton.solvers import SolverVBD +CONTACT_KE = 1.0e4 +CONTACT_KD = 1.0e-2 + SHAPE_CFG = newton.ModelBuilder.ShapeConfig( mu=0.1, - ke=1.0e6, - kd=1.0e3, + ke=CONTACT_KE, + kd=CONTACT_KD, gap=0.002, density=1.0e6, mu_torsional=0.0, @@ -45,8 +48,6 @@ CABLE_KINEMATIC_COUNT = 4 # first N rod bodies are inside the plug and follow it # Contact parameters for cable and ground plane (tuned for VBD). -CABLE_KE = 1.0e8 -CABLE_KD = 1.0e-3 CABLE_MU = 2.0 # Latch revolute-joint tuning. @@ -317,14 +318,12 @@ def __init__(self, viewer): radius=CABLE_RADIUS, cfg=dataclasses.replace( builder.default_shape_cfg, - ke=CABLE_KE, - kd=CABLE_KD, + ke=CONTACT_KE, + kd=CONTACT_KD, mu=CABLE_MU, ), bend_stiffness=1.0e-1, bend_damping=1.0e-1, - stretch_stiffness=1.0e9, - stretch_damping=1.0e-1, label="cable", ) @@ -384,10 +383,10 @@ def __init__(self, viewer): self.solver = SolverVBD( self.model, iterations=12, - friction_epsilon=0.1, - rigid_contact_k_start=1.0e5, rigid_body_contact_buffer_size=256, ) + for j in range(self.model.joint_count): + self.solver.set_joint_constraint_mode(j, False) self._rest_pos = plug_pos self.gizmo_tf = wp.transform(plug_pos, wp.quat_identity()) diff --git a/newton/tests/test_cable.py b/newton/tests/test_cable.py index c6e69c8a48..b5584f4f3d 100644 --- a/newton/tests/test_cable.py +++ b/newton/tests/test_cable.py @@ -171,6 +171,152 @@ def _drive_gripper_boxes_kernel( body_q[b] = wp.transform(pos, rot) +# ----------------------------------------------------------------------------- +# Device-side time kernels (for CUDA graph capture with kinematic bodies) +# ----------------------------------------------------------------------------- + + +@wp.kernel +def _advance_time(sim_time: wp.array[float], dt: float): + sim_time[0] = sim_time[0] + dt + + +@wp.kernel +def _set_kinematic_sinusoidal_pose( + body_id: wp.int32, + sim_time: wp.array[float], + anchor_z: float, + x_amp: float, + x_freq: float, + body_q: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], +): + t = wp.float32(sim_time[0]) + dx = x_amp * wp.sin(x_freq * t) + body_q[body_id] = wp.transform(wp.vec3(dx, 0.0, anchor_z), wp.quat_identity()) + body_qd[body_id] = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def _set_kinematic_sinusoidal_xy_pose( + body_id: wp.int32, + sim_time: wp.array[float], + anchor_z: float, + x_amp: float, + x_freq: float, + y_amp: float, + y_freq: float, + body_q: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], +): + t = wp.float32(sim_time[0]) + dx = x_amp * wp.sin(x_freq * t) + dy = y_amp * wp.sin(y_freq * t) + body_q[body_id] = wp.transform(wp.vec3(dx, dy, anchor_z), wp.quat_identity()) + body_qd[body_id] = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def _set_kinematic_d6_pose( + body_id: wp.int32, + sim_time: wp.array[float], + anchor_z: float, + x_amp: float, + x_freq: float, + ang_amp: float, + ang_freq: float, + body_q: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], +): + t = wp.float32(sim_time[0]) + dx = x_amp * wp.sin(x_freq * t) + ang_y = ang_amp * wp.sin(ang_freq * t) + half = ang_y * 0.5 + q_anchor = wp.quat(0.0, wp.sin(half), 0.0, wp.cos(half)) + body_q[body_id] = wp.transform(wp.vec3(dx, 0.0, anchor_z), q_anchor) + body_qd[body_id] = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def _set_kinematic_linear_rotating_pose( + body_id: wp.int32, + sim_time: wp.array[float], + anchor_z: float, + velocity_x: float, + angular_velocity_z: float, + body_q: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], +): + t = wp.float32(sim_time[0]) + x_kin = velocity_x * t + angle_z = angular_velocity_z * t + half = angle_z * 0.5 + q_kin = wp.quat(0.0, 0.0, wp.sin(half), wp.cos(half)) + body_q[body_id] = wp.transform(wp.vec3(x_kin, 0.0, anchor_z), q_kin) + body_qd[body_id] = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def _drive_gripper_boxes_graph_kernel( + ramp_time: float, + sim_time: wp.array[float], + body_ids: wp.array[wp.int32], + signs: wp.array[wp.float32], + anchor_p: wp.vec3, + anchor_q: wp.quat, + seg_half_len: float, + target_offset_mag: float, + initial_offset_mag: float, + pull_start_time: float, + pull_ramp_time: float, + pull_distance: float, + body_q: wp.array[wp.transform], +): + """Graph-capturable variant of _drive_gripper_boxes_kernel that reads time from a device array.""" + tid = wp.tid() + b = body_ids[tid] + sgn = signs[tid] + rot = anchor_q + center = anchor_p + wp.quat_rotate(rot, wp.vec3(0.0, 0.0, seg_half_len)) + t = wp.float32(sim_time[0]) + pull_end_time = wp.float32(pull_start_time + pull_ramp_time) + t_eff = wp.min(t, pull_end_time) + u = wp.clamp(t_eff / wp.float32(ramp_time), 0.0, 1.0) + offset_mag = (1.0 - u) * initial_offset_mag + u * target_offset_mag + tp = wp.clamp((t_eff - wp.float32(pull_start_time)) / wp.float32(pull_ramp_time), 0.0, 1.0) + pull = wp.float32(pull_distance) * tp + pull_dir = wp.quat_rotate(rot, wp.vec3(0.0, 0.0, 1.0)) + local_off = wp.vec3(0.0, sgn * offset_mag, 0.0) + pos = center + pull_dir * pull + wp.quat_rotate(rot, local_off) + body_q[b] = wp.transform(pos, rot) + + +# ----------------------------------------------------------------------------- +# Graph-capture helper +# ----------------------------------------------------------------------------- + + +def _run_sim_loop(simulate_fn, num_steps, device): + """Run a simulation loop with optional CUDA graph capture. + + ``simulate_fn()`` must be graph-capturable: no host-side branching, no + scalar time arguments — use device-side ``sim_time`` arrays and the + ``_advance_time`` kernel instead. + """ + use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + graph = None + if use_cuda_graph: + with wp.ScopedCapture(device) as capture: + simulate_fn() + graph = capture.graph + + for _ in range(num_steps): + if graph is not None: + wp.capture_launch(graph) + else: + simulate_fn() + + # ----------------------------------------------------------------------------- # Geometry helpers # ----------------------------------------------------------------------------- @@ -235,8 +381,8 @@ def _build_cable_chain( """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Geometry: straight cable along +X, centered around the origin @@ -250,8 +396,6 @@ def _build_cable_chain( radius=0.05, bend_stiffness=bend_stiffness, bend_damping=bend_damping, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, label="test_cable_chain", ) @@ -277,8 +421,8 @@ def _build_cable_loop(device, num_links: int = 6): """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Geometry: points on a circle in the X-Y plane at fixed height @@ -302,8 +446,6 @@ def _build_cable_loop(device, num_links: int = 6): radius=0.05, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, closed=True, label="test_cable_loop", ) @@ -629,8 +771,8 @@ def _cable_bend_stiffness_impl(test: unittest.TestCase, device): num_links = 10 builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Place cables far apart along Y so they don't interact. @@ -648,8 +790,6 @@ def _cable_bend_stiffness_impl(test: unittest.TestCase, device): radius=0.05, bend_stiffness=k, bend_damping=1.0e1, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, label=f"bend_stiffness_{k:.0e}", ) @@ -675,13 +815,16 @@ def _cable_bend_stiffness_impl(test: unittest.TestCase, device): num_steps = 20 # Run for a short duration to let bending respond to gravity - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): state0.clear_forces() model.collide(state0, contacts) solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + final_q = state0.body_q.numpy() tip_heights = np.array([final_q[tip_body, 2] for tip_body in tip_bodies], dtype=float) @@ -731,13 +874,16 @@ def _cable_sagging_and_stability_impl(test: unittest.TestCase, device): initial_q = state0.body_q.numpy().copy() z_initial = initial_q[:, 2] - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): state0.clear_forces() model.collide(state0, contacts) solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + final_q = state0.body_q.numpy() z_final = final_q[:, 2] @@ -765,8 +911,8 @@ def _cable_twist_response_impl(test: unittest.TestCase, device): # This isolates twist response when rotating the first (anchored) capsule about its local axis. builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 z_height = 3.0 @@ -790,8 +936,6 @@ def _cable_twist_response_impl(test: unittest.TestCase, device): radius=0.05, bend_stiffness=1.0e4, bend_damping=0.0, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, label="twist_chain_orthogonal", ) @@ -837,13 +981,16 @@ def _cable_twist_response_impl(test: unittest.TestCase, device): num_steps = 20 # Run a short simulation to let twist propagate - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): state0.clear_forces() model.collide(state0, contacts) solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + final_q = state0.body_q.numpy() # Check capsule attachments remain good @@ -931,8 +1078,8 @@ def _two_layer_cable_pile_collision_impl(test: unittest.TestCase, device): builder = newton.ModelBuilder() # Contact material (stiff contacts, noticeable friction) - builder.default_shape_cfg.ke = 1.0e5 - builder.default_shape_cfg.kd = 1.0e-1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Cable geometric parameters @@ -998,8 +1145,6 @@ def _two_layer_cable_pile_collision_impl(test: unittest.TestCase, device): radius=cable_radius, bend_stiffness=bend_stiffness, bend_damping=1.0e-1, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, label=f"pile_l{layer}_{lane}", ) cable_bodies.extend(rod_bodies) @@ -1020,13 +1165,17 @@ def _two_layer_cable_pile_collision_impl(test: unittest.TestCase, device): # Let the pile settle under gravity and contact num_steps = 20 - for _step in range(num_steps): + + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): state0.clear_forces() model.collide(state0, contacts) solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + body_q = state0.body_q.numpy() positions = body_q[:, :3] z_positions = positions[:, 2] @@ -1095,8 +1244,8 @@ def _two_layer_cable_pile_collision_impl(test: unittest.TestCase, device): def _cable_ball_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, device): """Cable VBD: BALL joint should keep rod start endpoint attached to a kinematic anchor.""" builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Kinematic anchor body at the rod start point. @@ -1134,8 +1283,6 @@ def _cable_ball_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, device radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_ball_joint_attach", ) @@ -1171,25 +1318,36 @@ def _cable_ball_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, device sim_dt = frame_dt / sim_substeps num_steps = 20 - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, 0.0, anchor_pos[2]), wp.quat_identity()) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_sinusoidal_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) - err = _compute_ball_joint_anchor_error(model, state0.body_q, j_ball) - test.assertLess(err, 1.0e-3) + _run_sim_loop(simulate, num_steps, device) + + err = _compute_ball_joint_anchor_error(model, state0.body_q, j_ball) + test.assertLess(err, 1.0e-3, f"BALL joint: final anchor error {err:.6f} m > 1e-3 m") # Also verify the rod joints remained well-attached along the chain. final_q = state0.body_q.numpy() @@ -1239,8 +1397,8 @@ def _cable_fixed_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic cable along one axis can't fully demonstrate this because gravity only tests one orientation. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -1275,8 +1433,6 @@ def _cable_fixed_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_fixed_joint_attach_x", ) @@ -1303,8 +1459,6 @@ def _cable_fixed_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_fixed_joint_attach_y", ) @@ -1337,30 +1491,41 @@ def _cable_fixed_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic sim_dt = frame_dt / sim_substeps num_steps = 20 - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, 0.0, anchor_pos[2]), wp.quat_identity()) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_sinusoidal_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) + + _run_sim_loop(simulate, num_steps, device) - pos_err_x, ang_err_x = _compute_fixed_joint_frame_error(model, state0.body_q, j_fixed_x) - test.assertLess(pos_err_x, 1.0e-3) - test.assertLess(ang_err_x, 2.0e-2) + pos_err_x, ang_err_x = _compute_fixed_joint_frame_error(model, state0.body_q, j_fixed_x) + test.assertLess(pos_err_x, 1.0e-3, f"FIXED joint (X): pos error {pos_err_x:.6f}") + test.assertLess(ang_err_x, 2.0e-2, f"FIXED joint (X): ang error {ang_err_x:.4f}") - pos_err_y, ang_err_y = _compute_fixed_joint_frame_error(model, state0.body_q, j_fixed_y) - test.assertLess(pos_err_y, 1.0e-3) - test.assertLess(ang_err_y, 2.0e-2) + pos_err_y, ang_err_y = _compute_fixed_joint_frame_error(model, state0.body_q, j_fixed_y) + test.assertLess(pos_err_y, 1.0e-3, f"FIXED joint (Y): pos error {pos_err_y:.6f}") + test.assertLess(ang_err_y, 2.0e-2, f"FIXED joint (Y): ang error {ang_err_y:.4f}") final_q = state0.body_q.numpy() test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms detected in FIXED joint test") @@ -1419,8 +1584,8 @@ def _cable_revolute_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, de so it sags. For cable Y, gravity creates torque about X (constrained), so it stays rigid. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -1455,8 +1620,6 @@ def _cable_revolute_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, de radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_revolute_joint_attach_x", ) @@ -1486,8 +1649,6 @@ def _cable_revolute_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, de radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_revolute_joint_attach_y", ) @@ -1524,30 +1685,41 @@ def _cable_revolute_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, de sim_dt = frame_dt / sim_substeps num_steps = 20 - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, 0.0, anchor_pos[2]), wp.quat_identity()) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_sinusoidal_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) + + _run_sim_loop(simulate, num_steps, device) - pos_err_x, ang_perp_err_x, _rot_free_x = _compute_revolute_joint_error(model, state0.body_q, j_revolute_x) - test.assertLess(pos_err_x, 1.0e-3) - test.assertLess(ang_perp_err_x, 2.0e-2) + pos_err_x, ang_perp_err_x, _ = _compute_revolute_joint_error(model, state0.body_q, j_revolute_x) + test.assertLess(pos_err_x, 1.0e-3, f"REVOLUTE joint (X): pos error {pos_err_x:.6f}") + test.assertLess(ang_perp_err_x, 2.0e-2, f"REVOLUTE joint (X): ang perp error {ang_perp_err_x:.4f}") - pos_err_y, ang_perp_err_y, _rot_free_y = _compute_revolute_joint_error(model, state0.body_q, j_revolute_y) - test.assertLess(pos_err_y, 1.0e-3) - test.assertLess(ang_perp_err_y, 2.0e-2) + pos_err_y, ang_perp_err_y, _ = _compute_revolute_joint_error(model, state0.body_q, j_revolute_y) + test.assertLess(pos_err_y, 1.0e-3, f"REVOLUTE joint (Y): pos error {pos_err_y:.6f}") + test.assertLess(ang_perp_err_y, 2.0e-2, f"REVOLUTE joint (Y): ang perp error {ang_perp_err_y:.4f}") final_q = state0.body_q.numpy() test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms detected in REVOLUTE joint test") @@ -1614,8 +1786,8 @@ def _cable_revolute_drive_tracks_target_impl(test: unittest.TestCase, device): toward it despite gravity. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -1647,8 +1819,6 @@ def _cable_revolute_drive_tracks_target_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_revolute_drive", ) @@ -1699,12 +1869,15 @@ def _cable_revolute_drive_tracks_target_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 30 - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + # Joint constraint checks. pos_err, ang_perp_err, rot_free = _compute_revolute_joint_error(model, state0.body_q, rev_idx) test.assertLess(pos_err, 1.0e-3, f"Revolute drive: position error {pos_err:.6f}") @@ -1738,8 +1911,8 @@ def _cable_revolute_drive_limit_impl(test: unittest.TestCase, device): the cable should reach the limit bound, not the drive target. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -1770,8 +1943,6 @@ def _cable_revolute_drive_limit_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_revolute_drive_limit", ) @@ -1825,12 +1996,15 @@ def _cable_revolute_drive_limit_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 30 - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + # Joint constraint checks. pos_err, ang_perp_err, rot_free = _compute_revolute_joint_error(model, state0.body_q, rev_idx) test.assertLess(pos_err, 2.0e-3, f"Revolute drive limit: position error {pos_err:.6f}") @@ -1872,8 +2046,8 @@ def _cable_prismatic_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, d along the free axis is a degenerate configuration (it can slide away from the anchor). """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -1906,8 +2080,6 @@ def _cable_prismatic_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, d radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_prismatic_joint_attach", ) @@ -1945,27 +2117,39 @@ def _cable_prismatic_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, d sim_dt = frame_dt / sim_substeps num_steps = 20 - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) - dy = wp.float32(0.04 * np.sin(2.0 * t)) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, dy, anchor_pos[2]), wp.quat_identity()) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_sinusoidal_xy_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + 0.04, + 2.0, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) + + _run_sim_loop(simulate, num_steps, device) - pos_perp_err, ang_err, _c_along = _compute_prismatic_joint_error(model, state0.body_q, j_prismatic) - test.assertLess(pos_perp_err, 1.0e-3) - test.assertLess(ang_err, 2.0e-2) + pos_perp_err, ang_err, _c_along = _compute_prismatic_joint_error(model, state0.body_q, j_prismatic) + test.assertLess(pos_perp_err, 1.0e-3, "PRISMATIC joint: perpendicular position error too large") + test.assertLess(ang_err, 2.0e-2, "PRISMATIC joint: angular error too large") final_q = state0.body_q.numpy() test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms detected in PRISMATIC joint test") @@ -2002,8 +2186,8 @@ def _cable_prismatic_drive_tracks_target_impl(test: unittest.TestCase, device): converge toward it despite gravity. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2035,8 +2219,6 @@ def _cable_prismatic_drive_tracks_target_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_prismatic_drive", ) @@ -2087,12 +2269,15 @@ def _cable_prismatic_drive_tracks_target_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 30 - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + # Joint constraint checks. pos_perp_err, ang_err, c_along = _compute_prismatic_joint_error(model, state0.body_q, prismatic_idx) test.assertLess(pos_perp_err, 2.0e-3, f"Prismatic drive: perpendicular position error {pos_perp_err:.6f}") @@ -2126,8 +2311,8 @@ def _cable_prismatic_drive_limit_impl(test: unittest.TestCase, device): the cable should reach the limit bound, not the drive target. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2158,8 +2343,6 @@ def _cable_prismatic_drive_limit_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_prismatic_drive_limit", ) @@ -2213,12 +2396,15 @@ def _cable_prismatic_drive_limit_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 30 - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + # Joint constraint checks. pos_perp_err, ang_err, c_along = _compute_prismatic_joint_error(model, state0.body_q, prismatic_idx) test.assertLess(pos_perp_err, 2.0e-3, f"Prismatic drive limit: perp pos error {pos_perp_err:.6f}") @@ -2260,12 +2446,12 @@ def _cable_d6_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, device): Anchor oscillates in X and rotates around Y. The free linear axis allows the cable to slide (not follow the X motion). Gravity in -Z stresses the locked Z linear constraint. The anchor rotation around Y - directly exercises the free angular Y axis — the cable should not + directly exercises the free angular Y axis - the cable should not follow the rotation. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2299,8 +2485,6 @@ def _cable_d6_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_d6_joint_attach", ) @@ -2335,29 +2519,39 @@ def _cable_d6_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 20 - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) - # Rotate anchor around Y to exercise the free angular Y DOF - ang_y = wp.float32(0.2 * np.sin(2.0 * t)) - q_anchor = wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), ang_y) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, 0.0, anchor_pos[2]), q_anchor) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_d6_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + 0.2, + 2.0, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) + + _run_sim_loop(simulate, num_steps, device) - pos_perp_err, ang_perp_err, _d_along, _rot_free = _compute_d6_joint_error(model, state0.body_q, j_d6) - test.assertLess(pos_perp_err, 1.0e-3) - test.assertLess(ang_perp_err, 2.0e-2) + pos_perp_err, ang_perp_err, _d_along, _rot_free = _compute_d6_joint_error(model, state0.body_q, j_d6) + test.assertLess(pos_perp_err, 1.0e-3, "D6 joint: perpendicular position error too large") + test.assertLess(ang_perp_err, 2.0e-2, "D6 joint: perpendicular angular error too large") final_q = state0.body_q.numpy() test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms in D6 joint test") @@ -2399,8 +2593,8 @@ def _cable_d6_joint_all_locked_impl(test: unittest.TestCase, device): follow exactly, matching the lock_xyz config in example_cable_d6_joints. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2432,8 +2626,6 @@ def _cable_d6_joint_all_locked_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_d6_all_locked", ) @@ -2468,26 +2660,37 @@ def _cable_d6_joint_all_locked_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 20 - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, 0.0, anchor_pos[2]), wp.quat_identity()) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_sinusoidal_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) - pos_err, ang_err = _compute_fixed_joint_frame_error(model, state0.body_q, j_d6) - test.assertLess(pos_err, 1.0e-3) - test.assertLess(ang_err, 2.0e-2) + _run_sim_loop(simulate, num_steps, device) + + pos_err, ang_err = _compute_fixed_joint_frame_error(model, state0.body_q, j_d6) + test.assertLess(pos_err, 1.0e-3, "D6 all-locked: position error too large") + test.assertLess(ang_err, 2.0e-2, "D6 all-locked: angular error too large") final_q = state0.body_q.numpy() test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms in D6 all-locked test") @@ -2518,8 +2721,8 @@ def _cable_d6_joint_locked_x_impl(test: unittest.TestCase, device): Matches the lock_x config in example_cable_d6_joints. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2553,8 +2756,6 @@ def _cable_d6_joint_locked_x_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_d6_locked_x", ) @@ -2593,39 +2794,50 @@ def _cable_d6_joint_locked_x_impl(test: unittest.TestCase, device): locked_axis_local = wp.vec3(1.0, 0.0, 0.0) - for _step in range(num_steps): - for _substep in range(sim_substeps): - t = (_step * sim_substeps + _substep) * sim_dt - dx = wp.float32(0.05 * np.sin(1.5 * t)) + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) - pose = wp.transform(wp.vec3(dx, 0.0, anchor_pos[2]), wp.quat_identity()) + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): wp.launch( - _set_kinematic_body_pose, + _set_kinematic_sinusoidal_pose, dim=1, - inputs=[wp.int32(anchor), pose, state0.body_q, state0.body_qd], + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + 0.05, + 1.5, + state0.body_q, + state0.body_qd, + ], device=device, ) - model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) + + _run_sim_loop(simulate, num_steps, device) + + # Position error along the locked X axis. + X_wp, X_wc = _get_joint_world_frames(model, state0.body_q, j_d6) + x_p = wp.transform_get_translation(X_wp) + x_c = wp.transform_get_translation(X_wc) + q_wp = wp.transform_get_rotation(X_wp) + axis_world = wp.normalize(wp.quat_rotate(q_wp, locked_axis_local)) + d_locked = abs(float(wp.dot(x_c - x_p, axis_world))) + test.assertLess(d_locked, 1.0e-3, "D6 locked X: position error along locked axis") - # Position error along the locked X axis. - X_wp, X_wc = _get_joint_world_frames(model, state0.body_q, j_d6) - x_p = wp.transform_get_translation(X_wp) - x_c = wp.transform_get_translation(X_wc) - q_wp = wp.transform_get_rotation(X_wp) - axis_world = wp.normalize(wp.quat_rotate(q_wp, locked_axis_local)) - d_locked = abs(float(wp.dot(x_c - x_p, axis_world))) - test.assertLess(d_locked, 1.0e-3, "D6 locked X: position error along locked axis") - - # Angular error (all angular locked). - q_wc = wp.transform_get_rotation(X_wc) - q_rel = wp.normalize(wp.mul(wp.quat_inverse(q_wp), q_wc)) - q_rest = _get_joint_rest_relative_rotation(model, j_d6) - q_err = wp.normalize(wp.mul(q_rel, wp.quat_inverse(q_rest))) - ang_err = float(2.0 * wp.acos(wp.clamp(wp.abs(q_err[3]), 0.0, 1.0))) - test.assertLess(ang_err, 2.0e-2, "D6 locked X: angular error") + # Angular error (all angular locked). + q_wc = wp.transform_get_rotation(X_wc) + q_rel = wp.normalize(wp.mul(wp.quat_inverse(q_wp), q_wc)) + q_rest = _get_joint_rest_relative_rotation(model, j_d6) + q_err = wp.normalize(wp.mul(q_rel, wp.quat_inverse(q_rest))) + ang_err = float(2.0 * wp.acos(wp.clamp(wp.abs(q_err[3]), 0.0, 1.0))) + test.assertLess(ang_err, 2.0e-2, "D6 locked X: angular error") final_q = state0.body_q.numpy() test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms in D6 locked-X test") @@ -2670,8 +2882,8 @@ def _cable_d6_drive_tracks_target_impl(test: unittest.TestCase, device): toward them despite gravity. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2703,8 +2915,6 @@ def _cable_d6_drive_tracks_target_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_d6_drive", ) @@ -2762,12 +2972,15 @@ def _cable_d6_drive_tracks_target_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 30 - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + # Joint constraint checks. pos_perp_err, ang_perp_err, c_along, rot_free = _compute_d6_joint_error(model, state0.body_q, d6_idx) test.assertLess(pos_perp_err, 2.0e-3, f"D6 drive: perpendicular position error {pos_perp_err:.6f}") @@ -2810,8 +3023,8 @@ def _cable_d6_drive_limit_impl(test: unittest.TestCase, device): cable should reach the limit bounds, not the drive targets. """ builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 anchor_pos = wp.vec3(0.0, 0.0, 3.0) @@ -2842,8 +3055,6 @@ def _cable_d6_drive_limit_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label="test_cable_d6_drive_limit", ) @@ -2916,12 +3127,15 @@ def _cable_d6_drive_limit_impl(test: unittest.TestCase, device): sim_dt = frame_dt / sim_substeps num_steps = 30 - for _step in range(num_steps): + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + # Locked DOF checks. pos_perp_err, ang_perp_err, c_along, rot_free = _compute_d6_joint_error(model, state0.body_q, d6_idx) test.assertLess(pos_perp_err, 2.0e-3, f"D6 drive limit: perp pos error {pos_perp_err:.6f}") @@ -2979,6 +3193,7 @@ def _cable_kinematic_gripper_picks_capsule_impl(test: unittest.TestCase, device) # Contact/friction: large mu to encourage sticking if kinematic friction is working. builder.default_shape_cfg.mu = 1.0e3 + builder.default_shape_cfg.ke = 1.0e4 # Payload: capsule sized to match old box AABB (0.20, 0.10, 0.10) in (X,Y,Z) box_hx = 0.10 @@ -3068,7 +3283,7 @@ def _cable_kinematic_gripper_picks_capsule_impl(test: unittest.TestCase, device) fps = 60.0 frame_dt = 1.0 / fps - sim_substeps = 1 + sim_substeps = 2 sim_dt = frame_dt / sim_substeps # Record initial pose @@ -3077,17 +3292,20 @@ def _cable_kinematic_gripper_picks_capsule_impl(test: unittest.TestCase, device) # Run a fixed number of frames for a lightweight regression test. num_frames = 100 - sim_time = 0.0 num_steps = num_frames * sim_substeps - for _step in range(num_steps): + + sim_time_arr = wp.zeros(1, dtype=float, device=device) + + def simulate(): + nonlocal state0, state1 state0.clear_forces() wp.launch( - kernel=_drive_gripper_boxes_kernel, + kernel=_drive_gripper_boxes_graph_kernel, dim=2, inputs=[ float(ramp_time), - float(sim_time), + sim_time_arr, gripper_body_ids, gripper_signs, anchor_p, @@ -3106,8 +3324,9 @@ def _cable_kinematic_gripper_picks_capsule_impl(test: unittest.TestCase, device) model.collide(state0, contacts) solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) - sim_time += sim_dt + _run_sim_loop(simulate, num_steps, device) qf = state0.body_q.numpy() test.assertTrue(np.isfinite(qf).all(), "Non-finite body transforms detected in gripper friction test") @@ -3134,8 +3353,8 @@ def _cable_kinematic_gripper_picks_capsule_impl(test: unittest.TestCase, device) def _cable_graph_y_junction_spanning_tree_impl(test: unittest.TestCase, device): """Cable graph: Y-junction should build (and simulate) with wrap_in_articulation=True.""" builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e5 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Simple Y: 0-1-2 and 1-3 @@ -3158,8 +3377,6 @@ def _cable_graph_y_junction_spanning_tree_impl(test: unittest.TestCase, device): cfg=builder.default_shape_cfg.copy(), bend_stiffness=1.0e2, bend_damping=1.0e-2, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, label="ut_cable_graph_y", wrap_in_articulation=True, ) @@ -3228,18 +3445,22 @@ def union_uf(a: int, b: int): z_init_min = float(np.min(q_init[rod_bodies, 2])) frame_dt = 1.0 / 60.0 - sim_substeps = 5 + sim_substeps = 6 sim_dt = frame_dt / sim_substeps num_steps = 20 contacts = model.contacts() - for _step in range(num_steps): + + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): model.collide(state0, contacts) state0.clear_forces() solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + qf = state0.body_q.numpy() test.assertTrue(np.isfinite(qf).all(), "Non-finite body transforms detected in Y-junction graph simulation") @@ -3251,8 +3472,8 @@ def union_uf(a: int, b: int): def _cable_rod_ring_closed_in_articulation_impl(test: unittest.TestCase, device): """Closed ring via add_rod(closed=True) should build and simulate.""" builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 # Build a planar ring polyline (duplicate last point so the last segment returns to the start). @@ -3276,8 +3497,6 @@ def _cable_rod_ring_closed_in_articulation_impl(test: unittest.TestCase, device) cfg=builder.default_shape_cfg.copy(), bend_stiffness=1.0e2, bend_damping=1.0e-2, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, closed=True, label="ut_cable_rod_ring_closed", wrap_in_articulation=True, @@ -3308,7 +3527,7 @@ def _cable_rod_ring_closed_in_articulation_impl(test: unittest.TestCase, device) solver = newton.solvers.SolverVBD(model, iterations=10) frame_dt = 1.0 / 60.0 - sim_substeps = 5 + sim_substeps = 6 sim_dt = frame_dt / sim_substeps num_steps = 20 @@ -3316,13 +3535,17 @@ def _cable_rod_ring_closed_in_articulation_impl(test: unittest.TestCase, device) z_init_min = float(np.min(q_init[rod_bodies, 2])) contacts = model.contacts() - for _step in range(num_steps): + + def simulate(): + nonlocal state0, state1 for _substep in range(sim_substeps): state0.clear_forces() model.collide(state0, contacts) solver.step(state0, state1, control, contacts, sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + qf = state0.body_q.numpy() test.assertTrue(np.isfinite(qf).all(), "Non-finite body transforms detected in closed-ring simulation") @@ -3350,8 +3573,6 @@ def _cable_graph_default_quat_aligns_z_impl(test: unittest.TestCase, device): cfg=builder.default_shape_cfg.copy(), bend_stiffness=0.0, bend_damping=0.0, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, label="ut_cable_graph_quat", wrap_in_articulation=True, quaternions=None, @@ -3409,8 +3630,6 @@ def assert_body_pair_filtered(builder: newton.ModelBuilder, model: newton.Model, cfg=builder.default_shape_cfg.copy(), bend_stiffness=0.0, bend_damping=0.0, - stretch_stiffness=1.0e6, - stretch_damping=0.0, label="ut_cable_graph_y_filter", wrap_in_articulation=True, ) @@ -3446,11 +3665,11 @@ def assert_body_pair_filtered(builder: newton.ModelBuilder, model: newton.Model, def _collect_rigid_body_contact_forces_impl(test: unittest.TestCase, device): """VBD rigid contact-force query returns valid per-contact buffers.""" builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e3 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 0.5 - # Two overlapping dynamic boxes to guarantee rigid-rigid contact generation. + # Two overlapping dynamic boxes - initial overlap guarantees contact. b0 = builder.add_body(xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()), mass=1.0, label="box0") b1 = builder.add_body(xform=wp.transform(wp.vec3(0.0, 0.0, 0.15), wp.quat_identity()), mass=1.0, label="box1") builder.add_shape_box(b0, hx=0.1, hy=0.1, hz=0.1) @@ -3461,16 +3680,21 @@ def _collect_rigid_body_contact_forces_impl(test: unittest.TestCase, device): model.set_gravity((0.0, 0.0, 0.0)) state0 = model.state() + state1 = model.state() contacts = model.contacts() - solver = newton.solvers.SolverVBD(model, iterations=1) + control = model.control() + solver = newton.solvers.SolverVBD(model, iterations=2) dt = 1.0 / 60.0 - # Build contacts on the current state and query them directly. - # This keeps the test focused on contact-force extraction, not on integration dynamics. + # Collide + step so ALM state (penalty_k, lambda) gets populated. model.collide(state0, contacts) + body_q_prev_snapshot = wp.clone(solver.body_q_prev) + solver.step(state0, state1, control, contacts, dt) - c_b0, c_b1, c_p0w, c_p1w, c_f_b1, c_count = solver.collect_rigid_contact_forces(state0, contacts, dt) + c_b0, c_b1, c_p0w, c_p1w, c_f_b1, c_count = solver.collect_rigid_contact_forces( + state1.body_q, body_q_prev_snapshot, contacts, dt + ) count = int(c_count.numpy()[0]) # Buffer lengths must match rigid contact capacity. @@ -3481,8 +3705,8 @@ def _collect_rigid_body_contact_forces_impl(test: unittest.TestCase, device): test.assertEqual(int(c_p1w.shape[0]), expected_len) test.assertEqual(int(c_f_b1.shape[0]), expected_len) - # We set up overlapping boxes, so at least one rigid contact should be queryable. - test.assertGreater(count, 0, msg="Expected at least one rigid-rigid contact") + # Two overlapping boxes, so at least one rigid contact should be queryable. + test.assertGreater(count, 0, msg="Expected at least one rigid contact") b0_np = c_b0.numpy() b1_np = c_b1.numpy() @@ -3525,8 +3749,8 @@ def _cable_world_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic for joint_label, joint_kind in joint_configs: builder = newton.ModelBuilder() - builder.default_shape_cfg.ke = 1.0e2 - builder.default_shape_cfg.kd = 1.0e1 + builder.default_shape_cfg.ke = 1.0e4 + builder.default_shape_cfg.kd = 0.0 builder.default_shape_cfg.mu = 1.0 points, edge_q = _make_straight_cable_along_x(num_elements, segment_length, z_height=z_height) @@ -3540,8 +3764,6 @@ def _cable_world_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic radius=rod_radius, bend_stiffness=1.0e-1, bend_damping=1.0e-2, - stretch_stiffness=1.0e9, - stretch_damping=0.0, wrap_in_articulation=False, label=f"test_cable_world_{joint_kind}", ) @@ -3604,12 +3826,15 @@ def _cable_world_joint_attaches_rod_endpoint_impl(test: unittest.TestCase, devic solver = newton.solvers.SolverVBD(model, iterations=10) - for _step in range(num_steps): + def simulate(_model=model, _solver=solver, _control=control, _contacts=contacts): + nonlocal state0, state1 for _substep in range(sim_substeps): - model.collide(state0, contacts) - solver.step(state0, state1, control, contacts, dt=sim_dt) + _model.collide(state0, _contacts) + _solver.step(state0, state1, _control, _contacts, dt=sim_dt) state0, state1 = state1, state0 + _run_sim_loop(simulate, num_steps, device) + final_q = state0.body_q.numpy() test.assertTrue( np.isfinite(final_q).all(), @@ -3713,13 +3938,11 @@ def _joint_enabled_toggle_impl(test: unittest.TestCase, device): radius=rod_radius, bend_stiffness=1.0e1, bend_damping=1.0e-2, - stretch_stiffness=1.0e6, - stretch_damping=1.0e-2, wrap_in_articulation=False, label="test_joint_enabled_cable", ) - # BALL joint: anchor sphere → first rod body. + # BALL joint: anchor sphere -> first rod body. parent_anchor_local = wp.vec3(0.0, 0.0, -attach_offset) child_anchor_local = wp.vec3(0.0, 0.0, 0.0) j = builder.add_joint_ball( @@ -3750,22 +3973,22 @@ def step_n(n): solver.step(state0, state1, control, contacts, dt=sim_dt) state0, state1 = state1, state0 - # Phase 1: joint enabled (default) — cable stays attached to anchor. + # Phase 1: joint enabled (default) - cable stays attached to anchor. step_n(10) err_connected = _compute_ball_joint_anchor_error(model, state0.body_q, j) test.assertLess(err_connected, 1.0e-3, f"Phase 1 (enabled): pos error {err_connected:.6f} m > 1e-3") - # Phase 2: disable joint — cable detaches and falls under gravity. + # Phase 2: disable joint - cable detaches and falls under gravity. enabled_np = model.joint_enabled.numpy() enabled_np[j] = False model.joint_enabled.assign(wp.array(enabled_np, dtype=bool, device=device)) step_n(10) err_disabled = _compute_ball_joint_anchor_error(model, state0.body_q, j) test.assertGreater( - err_disabled, 5.0e-3, f"Phase 2 (disabled): pos error {err_disabled:.6f} m — cable did not separate" + err_disabled, 5.0e-3, f"Phase 2 (disabled): pos error {err_disabled:.6f} m - cable did not separate" ) - # Phase 3: re-enable joint — solver pulls cable back toward anchor. + # Phase 3: re-enable joint - solver pulls cable back toward anchor. enabled_np[j] = True model.joint_enabled.assign(wp.array(enabled_np, dtype=bool, device=device)) step_n(10) @@ -3777,10 +4000,128 @@ def step_n(n): ) +def _cable_fixed_joint_tracks_moving_kinematic_impl(test: unittest.TestCase, device): + """Cable VBD: fixed joint tracks a translating-and-rotating kinematic body. + + A short cable is attached via a hard FIXED joint to a kinematic body that + translates along +X and rotates about Z. Verifies that both positional and + angular joint errors stay bounded every substep, exercising the linear and + angular C0 snapshot paths against a moving kinematic parent. + """ + builder = newton.ModelBuilder() + + anchor_pos = wp.vec3(0.0, 0.0, 1.0) + anchor = builder.add_body(xform=wp.transform(anchor_pos, wp.quat_identity())) + builder.add_shape_sphere(anchor, radius=0.05) + builder.body_mass[anchor] = 0.0 + builder.body_inv_mass[anchor] = 0.0 + builder.body_inertia[anchor] = wp.mat33(0.0) + builder.body_inv_inertia[anchor] = wp.mat33(0.0) + + num_elements = 3 + segment_length = 0.05 + rod_radius = 0.01 + attach_offset = wp.float32(0.05 + rod_radius) + + points, edge_q = _make_straight_cable_along_x(num_elements, segment_length, z_height=float(anchor_pos[2])) + parent_anchor_local = wp.vec3(attach_offset, 0.0, 0.0) + anchor_world_attach = anchor_pos + wp.vec3(float(attach_offset), 0.0, 0.0) + offset = anchor_world_attach - points[0] + points = [p + offset for p in points] + + rod_bodies, rod_joints = builder.add_rod( + positions=points, + quaternions=edge_q, + radius=rod_radius, + bend_stiffness=1.0e-1, + bend_damping=1.0e-2, + wrap_in_articulation=False, + label="test_kinematic_track", + ) + + j_fixed = builder.add_joint_fixed( + parent=anchor, + child=rod_bodies[0], + parent_xform=wp.transform(parent_anchor_local, edge_q[0]), + child_xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()), + ) + builder.add_articulation([*rod_joints, j_fixed]) + + builder.color() + model = builder.finalize(device=device) + model.set_gravity((0.0, 0.0, -9.81)) + + state0 = model.state() + state1 = model.state() + control = model.control() + contacts = model.contacts() + + solver = newton.solvers.SolverVBD(model, iterations=20) + + frame_dt = 1.0 / 60.0 + sim_substeps = 10 + sim_dt = frame_dt / sim_substeps + num_frames = 10 + velocity_x = 0.3 # m/s + angular_velocity_z = 1.0 # rad/s + + pos_tol = 1.5e-2 + ang_tol = 5.0e-2 + + sim_time_arr = wp.zeros(1, dtype=float, device=device) + anchor_id = wp.int32(anchor) + anchor_z = float(anchor_pos[2]) + + def simulate(): + nonlocal state0, state1 + for _substep in range(sim_substeps): + wp.launch( + _set_kinematic_linear_rotating_pose, + dim=1, + inputs=[ + anchor_id, + sim_time_arr, + anchor_z, + velocity_x, + angular_velocity_z, + state0.body_q, + state0.body_qd, + ], + device=device, + ) + model.collide(state0, contacts) + solver.step(state0, state1, control, contacts, dt=sim_dt) + state0, state1 = state1, state0 + wp.launch(_advance_time, dim=1, inputs=[sim_time_arr, sim_dt], device=device) + + _run_sim_loop(simulate, num_frames, device) + + pos_err, ang_err = _compute_fixed_joint_frame_error(model, state0.body_q, j_fixed) + test.assertLess( + pos_err, + pos_tol, + f"Fixed joint kinematic tracking: pos error {pos_err:.6f} m against moving kinematic body", + ) + test.assertLess( + ang_err, + ang_tol, + f"Fixed joint kinematic tracking: ang error {ang_err:.4f} rad against rotating kinematic body", + ) + + final_q = state0.body_q.numpy() + test.assertTrue(np.isfinite(final_q).all(), "Non-finite body transforms in kinematic tracking test") + + class TestCable(unittest.TestCase): pass +add_function_test( + TestCable, + "test_cable_fixed_joint_tracks_moving_kinematic", + _cable_fixed_joint_tracks_moving_kinematic_impl, + devices=devices, +) add_function_test( TestCable, "test_joint_enabled_toggle", diff --git a/newton/tests/test_collision_cloth.py b/newton/tests/test_collision_cloth.py index fc5f18548c..602ce71a9d 100644 --- a/newton/tests/test_collision_cloth.py +++ b/newton/tests/test_collision_cloth.py @@ -780,53 +780,6 @@ def test_mesh_ground_collision_index(test, device): test.assertTrue(np.allclose(normals[:, 2], 0.0, atol=1e-6)) -def test_avbd_particle_ground_penalty_grows(test, device): - """Regression: AVBD soft-contact penalty updates with particles + static ground only. - - When the model has particles and static shapes (shape_body == -1) but no rigid bodies - (body_count == 0), SolverVBD must still update the adaptive soft-contact penalty - (body-particle contact penalty k) so that particle-ground contacts do not remain - artificially soft. - """ - builder = newton.ModelBuilder(up_axis=newton.Axis.Z) - - # Ensure the contact stiffness cap is well above the initial k_start so we can observe growth. - builder.default_shape_cfg.ke = 1.0e9 - builder.default_shape_cfg.kd = 0.0 - - # Place a particle with positive penetration against the ground plane at z=0. - radius = 0.1 - builder.add_particle(pos=wp.vec3(0.0, 0.0, 0.0), vel=wp.vec3(0.0, 0.0, 0.0), mass=1.0, radius=radius) - builder.add_ground_plane() - builder.color() - - model = builder.finalize(device=device) - test.assertEqual(model.body_count, 0) - test.assertGreater(model.particle_count, 0) - test.assertGreater(model.shape_count, 0) - - vbd = SolverVBD(model, iterations=1, rigid_contact_k_start=1.0e2, rigid_avbd_beta=1.0e5) - - state_in = model.state() - state_out = model.state() - contacts = model.contacts() - model.collide(state_in, contacts) - - soft_count = int(contacts.soft_contact_count.numpy()[0]) - test.assertGreater(soft_count, 0) - - dt = 1.0 / 60.0 - control = model.control() - vbd._initialize_rigid_bodies(state_in, control, contacts, dt, update_rigid_history=True) - - k_before = float(vbd.body_particle_contact_penalty_k.numpy()[0]) - - vbd._solve_rigid_body_iteration(state_in, state_out, control, contacts, dt) - - k_after = float(vbd.body_particle_contact_penalty_k.numpy()[0]) - test.assertGreater(k_after, k_before) - - @wp.kernel def validate_vertex_collisions_distance_filter( max_query_radius: float, @@ -1117,9 +1070,6 @@ class TestCollision(unittest.TestCase): add_function_test(TestCollision, "test_edge_edge_collision", test_edge_edge_collision, devices=devices) add_function_test(TestCollision, "test_particle_collision", test_particle_collision, devices=devices) add_function_test(TestCollision, "test_mesh_ground_collision_index", test_mesh_ground_collision_index, devices=devices) -add_function_test( - TestCollision, "test_avbd_particle_ground_penalty_grows", test_avbd_particle_ground_penalty_grows, devices=devices -) add_function_test(TestCollision, "test_collision_filtering", test_collision_filtering, devices=devices) if __name__ == "__main__": diff --git a/newton/tests/test_kinematic_links.py b/newton/tests/test_kinematic_links.py index 62c2852ba5..5bb1b700be 100644 --- a/newton/tests/test_kinematic_links.py +++ b/newton/tests/test_kinematic_links.py @@ -329,7 +329,9 @@ def run_once(apply_force: bool): joint_qd[qd_start : qd_start + 6] = np.array([vx, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32) state_0.joint_q.assign(joint_q) state_0.joint_qd.assign(joint_qd) - newton.eval_fk(model, state_0.joint_q, state_0.joint_qd, state_0) + newton.eval_fk( + model, state_0.joint_q, state_0.joint_qd, state_0, body_flag_filter=newton.BodyFlags.KINEMATIC + ) state_0.clear_forces() if apply_force: