From f6f80cd0a05fd63903220c90b244f591983d2d55 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Fri, 10 Jul 2026 02:37:15 -0700 Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub. PiperOrigin-RevId: 945593266 Change-Id: Ic08842caabb92bc32974892ffae98e3295827094 --- .../mujoco_warp/_src/collision_convex.py | 141 +- .../mujoco_warp/_src/collision_core.py | 280 ++- .../mujoco_warp/_src/collision_driver.py | 47 +- .../mujoco_warp/_src/collision_flex.py | 1277 +++++++++-- .../mujoco_warp/_src/collision_gjk.py | 188 +- .../mujoco_warp/_src/constraint.py | 1864 ++++++++++++++--- .../mjx/third_party/mujoco_warp/_src/io.py | 202 +- .../third_party/mujoco_warp/_src/passive.py | 217 ++ .../third_party/mujoco_warp/_src/smooth.py | 120 +- .../third_party/mujoco_warp/_src/solver.py | 462 ++-- .../third_party/mujoco_warp/_src/support.py | 190 ++ .../third_party/mujoco_warp/_src/test_flex.py | 351 ++++ .../mjx/third_party/mujoco_warp/_src/types.py | 52 +- mjx/mujoco/mjx/warp/collision_driver.py | 21 +- mjx/mujoco/mjx/warp/forward.py | 582 ++--- mjx/mujoco/mjx/warp/types.py | 62 +- 16 files changed, 4878 insertions(+), 1178 deletions(-) create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py index 19d97952..1accd276 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py @@ -19,8 +19,11 @@ import warp as wp from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_margin_gap +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_material_params from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair_from_types from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import epa_phase @@ -707,8 +710,8 @@ def ccd_hfield_kernel_builder( return ccd_hfield_kernel -_CCD_OVERSUBSCRIBE_WAVES = 4 -_CCD_MIN_BLOCKS = 2 +_CCD_OVERSUBSCRIBE_WAVES = 2 +_CCD_MIN_BLOCKS = 8 @cache_kernel @@ -728,6 +731,17 @@ def ccd_kernel_builder( def eval_ccd_write_contact( # Model: opt_ccd_tolerance: wp.array[float], + geom_condim: wp.array[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_friction: wp.array2d[wp.vec3], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_friction: wp.array2d[vec5], # Data in: naconmax_in: int, naccdmax_in: int, @@ -756,13 +770,6 @@ def ccd_kernel_builder( nccd_in: wp.array[int], margin: float, gap: float, - condim: int, - friction: vec5, - solref: wp.vec2, - solreffriction: wp.vec2, - solimp: vec5, - x1: wp.vec3, - x2: wp.vec3, pairid: wp.vec2i, # Data out: contact_dist_out: wp.array[float], @@ -780,29 +787,26 @@ def ccd_kernel_builder( contact_type_out: wp.array[int], contact_geomcollisionid_out: wp.array[int], nacon_out: wp.array[int], - # Data out: overflow_out: wp.array[int], - ) -> int: - points = mat43() - witness1 = mat43() - witness2 = mat43() + ): geom1.margin = margin geom2.margin = margin + tolerance = opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]] is_collision_sensor = pairid[1] >= 0 if is_collision_sensor: cutoff = 1.0e32 else: cutoff = gap needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase( - opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], + tolerance, cutoff, gjk_iterations, geom1, geom2, geomtype1, geomtype2, - x1, - x2, + geom1.pos, + geom2.pos, ) ccdid = int(-1) @@ -814,9 +818,9 @@ def ccd_kernel_builder( if wp.static(warn_overflow): wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid) wp.atomic_or(overflow_out, worldid, OverflowType.CCD) - return 0 + return dist, ncollision, w1, w2, multiccd_idx = epa_phase( - opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], + tolerance, epa_iterations, gjk_result, geom1, @@ -831,8 +835,8 @@ def ccd_kernel_builder( epa_horizon_in[ccdid], ) - if dist >= gap and pairid[1] == -1: - return 0 + if dist >= gap and not is_collision_sensor: + return # CCD operates on margin-inflated shapes (support() inflates each geom by # 0.5 * margin). The returned dist is therefore relative to the inflated @@ -841,10 +845,16 @@ def ccd_kernel_builder( # with the primitive narrowphase, which reports un-inflated distances. dist += margin + witness1 = mat43() + witness2 = mat43() witness1[0] = w1 witness2[0] = w2 - if wp.static(use_multiccd or (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)): + if wp.static( + (use_multiccd or (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)) + and (geomtype1 == GeomType.BOX or geomtype1 == GeomType.MESH) + and (geomtype2 == GeomType.BOX or geomtype2 == GeomType.MESH) + ): if wp.static(geomtype1 == GeomType.MESH): # verify that geom1 mesh data is present for multicontact if geom1.mesh_polyadr < 0: @@ -879,23 +889,34 @@ def ccd_kernel_builder( geomtype2, ) - for i in range(ncollision): - points[i] = 0.5 * (witness1[i] + witness2[i]) - normal = witness1[0] - witness2[0] - frame = make_frame(normal) + condim, friction, solref, solreffriction, solimp = contact_material_params( + geom_condim, + geom_priority, + geom_solmix, + geom_solref, + geom_solimp, + geom_friction, + pair_dim, + pair_solref, + pair_solreffriction, + pair_solimp, + pair_friction, + geoms, + pairid[0], + worldid, + ) - # flip if collision sensor - if pairid[1] >= 0: + frame = make_frame(witness1[0] - witness2[0]) + if is_collision_sensor: frame *= -1.0 geoms = wp.vec2i(geoms[1], geoms[0]) - nactive = int(0) # number of contacts contributing to the physics for i in range(ncollision): - active = write_contact( + write_contact( naconmax_in, i, dist, - points[i], + 0.5 * (witness1[i] + witness2[i]), frame, margin, gap, @@ -923,9 +944,6 @@ def ccd_kernel_builder( contact_geomcollisionid_out, nacon_out, ) - nactive += active - - return nactive # runs convex collision on a set of geom pairs to recover contact info (non-heightfield) @wp.kernel(module="unique", enable_backward=False, launch_bounds=(block_dim, _CCD_MIN_BLOCKS)) @@ -1022,31 +1040,18 @@ def ccd_kernel_builder( continue worldid = collision_worldid_in[collisionid] - - _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( - geom_condim, - geom_priority, - geom_solmix, - geom_solref, - geom_solimp, - geom_friction, + pairid = collision_pairid_in[collisionid] + margin, gap = contact_margin_gap( geom_margin, geom_gap, - pair_dim, - pair_solref, - pair_solreffriction, - pair_solimp, pair_margin, pair_gap, - pair_friction, - collision_pair_in, - collision_pairid_in, - collisionid, + geoms, + pairid[0], worldid, ) - geom1, geom2 = geom_collision_pair( - geom_type, + geom1, geom2 = geom_collision_pair_from_types( geom_dataid, geom_size, mesh_vertadr, @@ -1065,12 +1070,25 @@ def ccd_kernel_builder( mesh_polymap, geom_xpos_in, geom_xmat_in, + geomtype1, + geomtype2, geoms, worldid, ) eval_ccd_write_contact( opt_ccd_tolerance, + geom_condim, + geom_priority, + geom_solmix, + geom_solref, + geom_solimp, + geom_friction, + pair_dim, + pair_solref, + pair_solreffriction, + pair_solimp, + pair_friction, naconmax_in, naccdmax_in, epa_vert_in, @@ -1097,14 +1115,7 @@ def ccd_kernel_builder( nccd_in, margin, gap, - condim, - friction, - solref, - solreffriction, - solimp, - geom1.pos, - geom2.pos, - collision_pairid_in[collisionid], + pairid, contact_dist_out, contact_pos_out, contact_frame_out, @@ -1126,11 +1137,15 @@ def ccd_kernel_builder( return ccd_kernel -def _ccd_grid_size(kernel, naconmax: int) -> int: +def _ccd_grid_size(kernel, naconmax: int, device) -> int: # Grid-stride launch width for the CCD kernel: a few device waves, capped at the contact # capacity. The kernel strides over the actual candidate count, so we avoid launching one # (mostly idle) thread per naconmax slot. - block_size, min_grid_size = wp.get_suggested_block_size(kernel) + if device.is_cpu: + # Warp forces CPU block_dim to 1 and has no CUDA occupancy information. + return naconmax + + block_size, min_grid_size = wp.get_suggested_block_size(kernel, device) return max(1, min(naconmax, _CCD_OVERSUBSCRIBE_WAVES * block_size * min_grid_size)) @@ -1326,7 +1341,7 @@ def convex_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table m.block_dim.convex_ccd, bool(m.opt.warn_overflow), ) - ccd_grid = _ccd_grid_size(ccd_k, d.naconmax) + ccd_grid = _ccd_grid_size(ccd_k, d.naconmax, d.ncollision.device) wp.launch( ccd_k, dim=ccd_grid, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py index dbf0fcec..0e906855 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py @@ -16,7 +16,7 @@ """Core collision types and utilities shared across collision modules.""" import dataclasses -from typing import Tuple +from typing import Any, Tuple import warp as wp @@ -63,9 +63,8 @@ class Geom: @wp.func -def geom_collision_pair( +def geom_collision_pair_from_types( # Model: - geom_type: wp.array[int], geom_dataid: wp.array2d[int], geom_size: wp.array2d[wp.vec3], mesh_vertadr: wp.array[int], @@ -86,6 +85,8 @@ def geom_collision_pair( geom_xpos_in: wp.array2d[wp.vec3], geom_xmat_in: wp.array2d[wp.mat33], # In: + geom_type1: int, + geom_type2: int, geoms: wp.vec2i, worldid: int, ) -> Tuple[Geom, Geom]: @@ -94,8 +95,6 @@ def geom_collision_pair( g1 = geoms[0] g2 = geoms[1] - geom_type1 = geom_type[g1] - geom_type2 = geom_type[g2] geom1.pos = geom_xpos_in[worldid, g1] geom1.rot = geom_xmat_in[worldid, g1] @@ -156,6 +155,61 @@ def geom_collision_pair( return geom1, geom2 +@wp.func +def geom_collision_pair( + # Model: + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_graph: wp.array[int], + mesh_polynum: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polymap: wp.array[int], + # Data in: + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + # In: + geoms: wp.vec2i, + worldid: int, +) -> Tuple[Geom, Geom]: + geom_type1 = geom_type[geoms[0]] + geom_type2 = geom_type[geoms[1]] + return geom_collision_pair_from_types( + geom_dataid, + geom_size, + mesh_vertadr, + mesh_vertnum, + mesh_graphadr, + mesh_vert, + mesh_graph, + mesh_polynum, + mesh_polyadr, + mesh_polynormal, + mesh_polyvertadr, + mesh_polyvertnum, + mesh_polyvert, + mesh_polymapadr, + mesh_polymapnum, + mesh_polymap, + geom_xpos_in, + geom_xmat_in, + geom_type1, + geom_type2, + geoms, + worldid, + ) + + @wp.func def write_contact( # Data in: @@ -234,7 +288,33 @@ def write_contact( @wp.func -def contact_params( +def contact_margin_gap( + # Model: + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + # In: + geoms: wp.vec2i, + pairid: int, + worldid: int, +) -> Tuple[float, float]: + if pairid > -1: + margin = pair_margin[worldid % pair_margin.shape[0], pairid] + gap = pair_gap[worldid % pair_gap.shape[0], pairid] + else: + g1 = geoms[0] + g2 = geoms[1] + margin_id = worldid % geom_margin.shape[0] + gap_id = worldid % geom_gap.shape[0] + margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2] + gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2] + + return margin, gap + + +@wp.func +def contact_material_params( # Model: geom_condim: wp.array[int], geom_priority: wp.array[int], @@ -242,35 +322,17 @@ def contact_params( geom_solref: wp.array2d[wp.vec2], geom_solimp: wp.array2d[vec5], geom_friction: wp.array2d[wp.vec3], - geom_margin: wp.array2d[float], - geom_gap: wp.array2d[float], pair_dim: wp.array[int], pair_solref: wp.array2d[wp.vec2], pair_solreffriction: wp.array2d[wp.vec2], pair_solimp: wp.array2d[vec5], - pair_margin: wp.array2d[float], - pair_gap: wp.array2d[float], pair_friction: wp.array2d[vec5], # In: - collision_pair_in: wp.array[wp.vec2i], - collision_pairid_in: wp.array[wp.vec2i], - cid: int, + geoms: wp.vec2i, + pairid: int, worldid: int, ): - """Resolve contact parameters for a collision pair. - - Uses explicit pair overrides when available, otherwise mixes geom-level - properties by priority and solmix weights. - """ - geoms = collision_pair_in[cid] - pairid = collision_pairid_in[cid][0] - - # TODO(team): early return if collision sensor but no contact - # (ie, pairid[0] < -1 and pairid[1] < 0) - if pairid > -1: - margin = pair_margin[worldid % pair_margin.shape[0], pairid] - gap = pair_gap[worldid % pair_gap.shape[0], pairid] condim = pair_dim[pairid] friction = pair_friction[worldid % pair_friction.shape[0], pairid] solref = pair_solref[worldid % pair_solref.shape[0], pairid] @@ -283,8 +345,6 @@ def contact_params( friction_id = worldid % geom_friction.shape[0] solref_id = worldid % geom_solref.shape[0] solimp_id = worldid % geom_solimp.shape[0] - margin_id = worldid % geom_margin.shape[0] - gap_id = worldid % geom_gap.shape[0] solmix1 = geom_solmix[solmix_id, g1] solmix2 = geom_solmix[solmix_id, g2] @@ -327,9 +387,6 @@ def contact_params( solreffriction = wp.vec2(0.0, 0.0) solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2] - # geom priority is ignored - margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2] - gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2] friction = vec5( wp.max(MJ_MINMU, friction[0]), @@ -339,6 +396,62 @@ def contact_params( wp.max(MJ_MINMU, friction[4]), ) + return condim, friction, solref, solreffriction, solimp + + +@wp.func +def contact_params( + # Model: + geom_condim: wp.array[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + pair_friction: wp.array2d[vec5], + # In: + collision_pair_in: wp.array[wp.vec2i], + collision_pairid_in: wp.array[wp.vec2i], + cid: int, + worldid: int, +): + """Resolve contact parameters for a collision pair. + + Uses explicit pair overrides when available, otherwise mixes geom-level + properties by priority and solmix weights. + """ + geoms = collision_pair_in[cid] + pairid = collision_pairid_in[cid][0] + + # TODO(team): early return if collision sensor but no contact + # (ie, pairid[0] < -1 and pairid[1] < 0) + + margin, gap = contact_margin_gap(geom_margin, geom_gap, pair_margin, pair_gap, geoms, pairid, worldid) + condim, friction, solref, solreffriction, solimp = contact_material_params( + geom_condim, + geom_priority, + geom_solmix, + geom_solref, + geom_solimp, + geom_friction, + pair_dim, + pair_solref, + pair_solreffriction, + pair_solimp, + pair_friction, + geoms, + pairid, + worldid, + ) + return geoms, margin, gap, condim, friction, solref, solreffriction, solimp @@ -357,6 +470,111 @@ class CollisionContext: collision_worldid: wp.array +@wp.func +def sap_binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) -> int: + """Binary search for the first element > value in sorted array.""" + while lower < upper: + mid = (lower + upper) >> 1 + if values[mid] > value: + upper = mid + else: + lower = mid + 1 + return upper + + +@wp.kernel +def sap_range( + # In: + n: int, + lower_in: wp.array2d[float], + upper_in: wp.array2d[float], + sort_index_in: wp.array2d[int], + # Out: + range_out: wp.array2d[int], +): + """Compute the sweep range for each sorted element.""" + worldid, sortedid = wp.tid() + + idx = sort_index_in[worldid, sortedid] + upper = upper_in[worldid, idx] + + limit = sap_binary_search(lower_in[worldid], upper, sortedid + 1, n) + limit = wp.min(n - 1, limit) + + range_out[worldid, sortedid] = limit - sortedid + + +@wp.kernel +def sap_sweep( + # In: + n: int, + sort_index_in: wp.array2d[int], + cumulative_sum_in: wp.array[int], + nsweep_in: int, + aabb_lower_in: wp.array2d[wp.vec3], + aabb_upper_in: wp.array2d[wp.vec3], + max_pairs: int, + # Out: + npairs_out: wp.array[int], + pair_id1_out: wp.array[int], + pair_id2_out: wp.array[int], + pair_worldid_out: wp.array[int], +): + """Generic SAP sweep: output AABB-overlapping pairs. + + This is the GPU equivalent of MuJoCo's mj_SAP function. It takes + axis-aligned bounding boxes and outputs pairs whose AABBs overlap + on all 3 axes. Domain-specific filtering is done by the caller. + """ + worldelemid = wp.tid() + + nworldelem = cumulative_sum_in.shape[0] + nworkpackages = cumulative_sum_in[nworldelem - 1] + + while worldelemid < nworkpackages: + # Binary search to find sortedid (i) and partner sortedid (j) + i = sap_binary_search(cumulative_sum_in, worldelemid, 0, nworldelem) + j = i + worldelemid + 1 + if i > 0: + j -= cumulative_sum_in[i - 1] + + worldid = i // n + i = i % n + j = j % n + + # Get actual element indices from sorted order + elem1 = sort_index_in[worldid, i] + elem2 = sort_index_in[worldid, j] + + # Ensure elem1 < elem2 for consistent ordering + if elem1 > elem2: + tmp = elem1 + elem1 = elem2 + elem2 = tmp + + worldelemid += nsweep_in + + # AABB overlap test on all 3 axes + lower1 = aabb_lower_in[worldid, elem1] + upper1 = aabb_upper_in[worldid, elem1] + lower2 = aabb_lower_in[worldid, elem2] + upper2 = aabb_upper_in[worldid, elem2] + + if lower1[0] > upper2[0] or lower2[0] > upper1[0]: + continue + if lower1[1] > upper2[1] or lower2[1] > upper1[1]: + continue + if lower1[2] > upper2[2] or lower2[2] > upper1[2]: + continue + + # Output this pair + idx = wp.atomic_add(npairs_out, 0, 1) + if idx < max_pairs: + pair_id1_out[idx] = elem1 + pair_id2_out[idx] = elem2 + pair_worldid_out[idx] = worldid + + def create_collision_context(naconmax: int) -> CollisionContext: """Create a CollisionContext with allocated arrays.""" return CollisionContext( diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py index 9b53f938..e65bab44 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py @@ -13,13 +13,15 @@ # limitations under the License. # ============================================================================== -from typing import Any, Optional +from typing import Optional import warp as wp from mujoco.mjx.third_party.mujoco_warp._src.collision_convex import convex_narrowphase from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext from mujoco.mjx.third_party.mujoco_warp._src.collision_core import create_collision_context +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import sap_binary_search +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import sap_range from mujoco.mjx.third_party.mujoco_warp._src.collision_flex import flex_collision from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import primitive_narrowphase from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf_narrowphase @@ -81,7 +83,6 @@ MJ_COLLISION_TABLE = { # TODO(team): Implement narrowphase flex collision support for: # - HFIELD -# - ELLIPSOID # - SDF MJ_FLEX_COLLISION_TABLE = { (GeomType.PLANE, GeomType.FLEX): CollisionType.PRIMITIVE, @@ -90,6 +91,7 @@ MJ_FLEX_COLLISION_TABLE = { (GeomType.BOX, GeomType.FLEX): CollisionType.PRIMITIVE, (GeomType.CYLINDER, GeomType.FLEX): CollisionType.PRIMITIVE, (GeomType.MESH, GeomType.FLEX): CollisionType.CONVEX, + (GeomType.ELLIPSOID, GeomType.FLEX): CollisionType.CONVEX, } @@ -369,18 +371,6 @@ def _add_geom_pair( collision_worldid_out[pairid] = worldid -@wp.func -def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) -> int: - while lower < upper: - mid = (lower + upper) >> 1 - if values[mid] > value: - upper = mid - else: - lower = mid + 1 - - return upper - - @cache_kernel def _sap_project(opt_broadphase: int): @wp.kernel(module="unique", enable_backward=False) @@ -430,31 +420,6 @@ def _sap_project(opt_broadphase: int): return sap_project -@wp.kernel -def _sap_range( - # Model: - ngeom: int, - # In: - projection_lower_in: wp.array2d[float], - projection_upper_in: wp.array2d[float], - sort_index_in: wp.array2d[int], - # Out: - range_out: wp.array2d[int], -): - worldid, geomid = wp.tid() - - # current bounding geom - idx = sort_index_in[worldid, geomid] - - upper = projection_upper_in[worldid, idx] - - limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom) - limit = wp.min(ngeom - 1, limit) - - # range of geoms for the sweep and prune process - range_out[worldid, geomid] = limit - geomid - - @cache_kernel def _sap_broadphase( opt_broadphase_filter: int, @@ -501,7 +466,7 @@ def _sap_broadphase( while worldgeomid < nworkpackages: # binary search to find current and next geom pair indices - i = _binary_search(cumulative_sum_in, worldgeomid, 0, nworldgeom) + i = sap_binary_search(cumulative_sum_in, worldgeomid, 0, nworldgeom) j = i + worldgeomid + 1 if i > 0: @@ -670,7 +635,7 @@ def sap_broadphase( ) wp.launch( - kernel=_sap_range, + kernel=sap_range, dim=(d.nworld, m.ngeom), inputs=[m.ngeom, projection_lower.reshape((-1, m.ngeom)), projection_upper, sort_index.reshape((-1, m.ngeom))], outputs=[range_], diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py index 0e2e3117..9e165b14 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py @@ -19,6 +19,8 @@ import warp as wp from mujoco.mjx.third_party.mujoco_warp._src import collision_primitive_core from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import sap_range +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import sap_sweep # TODO(team): consolidate _flex_sap_project with geom _sap_project from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES @@ -131,6 +133,7 @@ def _flex_triangle_geom_broadphase( and gtype != int(GeomType.BOX) and gtype != int(GeomType.CYLINDER) and gtype != int(GeomType.MESH) + and gtype != int(GeomType.ELLIPSOID) ): return @@ -190,6 +193,8 @@ def _flex_triangle_geom_broadphase( r_extent = wp.length(geom_half_size_local) elif gtype == int(GeomType.MESH): r_extent = wp.length(geom_half_size_local) + elif gtype == int(GeomType.ELLIPSOID): + r_extent = wp.length(geom_half_size_local) if wp.abs(signed_dist) > r_extent + margin + tri_radius: return @@ -353,6 +358,7 @@ def _flex_broadphase_plane( diff_center = flex_center - geom_pos dist_center = wp.dot(diff_center, plane_normal) + if dist_center - proj_half > margin: return @@ -415,6 +421,7 @@ def _write_flex_contact( contact_dim_out: wp.array[int], contact_geom_out: wp.array[wp.vec2i], contact_flex_out: wp.array[wp.vec2i], + contact_elem_out: wp.array[wp.vec2i], contact_vert_out: wp.array[wp.vec2i], contact_worldid_out: wp.array[int], contact_type_out: wp.array[int], @@ -452,8 +459,9 @@ def _write_flex_contact( contact_solimp_out[c_idx] = solimp contact_dim_out[c_idx] = condim contact_geom_out[c_idx] = wp.vec2i(geomid, -1) - contact_flex_out[c_idx] = wp.vec2i(flexid, elemid) - contact_vert_out[c_idx] = wp.vec2i(vertid, -1) + contact_flex_out[c_idx] = wp.vec2i(-1, flexid) + contact_elem_out[c_idx] = wp.vec2i(-1, elemid) + contact_vert_out[c_idx] = wp.vec2i(-1, vertid) contact_worldid_out[c_idx] = worldid contact_type_out[c_idx] = ContactType.CONSTRAINT contact_geomcollisionid_out[c_idx] = collisionid @@ -824,7 +832,7 @@ def _collide_mesh_triangle( geom2.margin = 0.0 geom2.index = -1 - centroid = (t1 + t2 + t3) * (1.0 / 3.0) + centroid = (t1 + t2 + t3) * wp.static(1.0 / 3.0) r_geom = wp.length(geom_size_val) d1 = wp.length(t1 - centroid) d2 = wp.length(t2 - centroid) @@ -1004,6 +1012,7 @@ def _flex_plane_narrowphase( contact_dim_out: wp.array[int], contact_geom_out: wp.array[wp.vec2i], contact_flex_out: wp.array[wp.vec2i], + contact_elem_out: wp.array[wp.vec2i], contact_vert_out: wp.array[wp.vec2i], contact_worldid_out: wp.array[int], contact_type_out: wp.array[int], @@ -1082,6 +1091,7 @@ def _flex_plane_narrowphase( contact_dim_out, contact_geom_out, contact_flex_out, + contact_elem_out, contact_vert_out, contact_worldid_out, contact_type_out, @@ -1671,6 +1681,380 @@ def _elements_overlap( return True +@wp.kernel +def _flex_sap_project( + # Model: + nflex: int, + flex_selfcollide: wp.array[int], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elem: wp.array[int], + flex_radius: wp.array[float], + flex_elemflexid: wp.array[int], + # Data in: + flexvert_xpos_in: wp.array2d[wp.vec3], + nworld_in: int, + # In: + nelem: int, + direction: wp.vec3, + # Out: + projection_lower_out: wp.array2d[float], + projection_upper_out: wp.array2d[float], + sort_index_out: wp.array2d[int], + elem_aabb_lower_out: wp.array2d[wp.vec3], + elem_aabb_upper_out: wp.array2d[wp.vec3], + segmented_index_out: wp.array[int], +): + worldid, elemid = wp.tid() + + flexid = flex_elemflexid[elemid] + + # Initialize sort index + sort_index_out[worldid, elemid] = elemid + + # Compute AABB from vertex positions + dim = flex_dim[flexid] + vert_adr = flex_vertadr[flexid] + elem_adr = flex_elemadr[flexid] + e = elemid - elem_adr + elem_data_idx = flex_elemdataadr[flexid] + e * (dim + 1) + + v = _get_element_vertices(flex_elem, dim, elem_data_idx) + + p0 = flexvert_xpos_in[worldid, vert_adr + v[0]] + p1 = flexvert_xpos_in[worldid, vert_adr + v[1]] + + aabb_min = wp.min(p0, p1) + aabb_max = wp.max(p0, p1) + + if dim >= 2: + p2 = flexvert_xpos_in[worldid, vert_adr + v[2]] + aabb_min = wp.min(aabb_min, p2) + aabb_max = wp.max(aabb_max, p2) + if dim >= 3: + p3 = flexvert_xpos_in[worldid, vert_adr + v[3]] + aabb_min = wp.min(aabb_min, p3) + aabb_max = wp.max(aabb_max, p3) + + radius = flex_radius[flexid] + rbound = 2.0 * radius + inflate = wp.vec3(rbound, rbound, rbound) + aabb_min = aabb_min - inflate + aabb_max = aabb_max + inflate + + elem_aabb_lower_out[worldid, elemid] = aabb_min + elem_aabb_upper_out[worldid, elemid] = aabb_max + + # Project AABB onto direction to get 1D interval + center = 0.5 * (aabb_min + aabb_max) + halfsize = 0.5 * (aabb_max - aabb_min) + proj_center = wp.dot(direction, center) + proj_radius = wp.abs(direction[0]) * halfsize[0] + wp.abs(direction[1]) * halfsize[1] + wp.abs(direction[2]) * halfsize[2] + + # If self-collision is disabled for this flex, push to infinity + if flex_selfcollide[flexid] == 0: + projection_lower_out[worldid, elemid] = MJ_MAXVAL + projection_upper_out[worldid, elemid] = MJ_MAXVAL + else: + projection_lower_out[worldid, elemid] = proj_center - proj_radius + projection_upper_out[worldid, elemid] = proj_center + proj_radius + + # Segmented sort boundaries + if elemid == 0: + segmented_index_out[worldid] = worldid * nelem + if worldid == nworld_in - 1: + segmented_index_out[nworld_in] = nworld_in * nelem + + +@wp.kernel +def _flex_sap_filter( + # Model: + flex_selfcollide: wp.array[int], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_elem: wp.array[int], + flex_elemflexid: wp.array[int], + # In: + raw_npairs_in: wp.array[int], + raw_pair_elem1_in: wp.array[int], + raw_pair_elem2_in: wp.array[int], + raw_pair_worldid_in: wp.array[int], + # Out: + npairs_out: wp.array[int], + pair_elem1_out: wp.array[int], + pair_elem2_out: wp.array[int], + pair_worldid_out: wp.array[int], +): + """Filter raw SAP pairs for flex self-collision.""" + tid = wp.tid() + + # Skip if beyond actual pair count + n_raw = raw_npairs_in[0] + if tid >= n_raw: + return + + elem1_global = raw_pair_elem1_in[tid] + elem2_global = raw_pair_elem2_in[tid] + + # Both elements must belong to the same flex + flexid1 = flex_elemflexid[elem1_global] + flexid2 = flex_elemflexid[elem2_global] + if flexid1 != flexid2: + return + + flexid = flexid1 + if flex_selfcollide[flexid] == 0: + return + + # Exclude elements sharing vertices/bodies + dim = flex_dim[flexid] + vert_adr = flex_vertadr[flexid] + elem_adr = flex_elemadr[flexid] + + e1 = elem1_global - elem_adr + e2 = elem2_global - elem_adr + elem_data_idx1 = flex_elemdataadr[flexid] + e1 * (dim + 1) + elem_data_idx2 = flex_elemdataadr[flexid] + e2 * (dim + 1) + v1_indices = _get_element_vertices(flex_elem, dim, elem_data_idx1) + v2_indices = _get_element_vertices(flex_elem, dim, elem_data_idx2) + + if _exclude_self_collision(flex_vertbodyid, v1_indices, dim + 1, v2_indices, dim + 1, vert_adr): + return + + # Output this pair + idx = wp.atomic_add(npairs_out, 0, 1) + if idx < pair_elem1_out.shape[0]: + pair_elem1_out[idx] = elem1_global + pair_elem2_out[idx] = elem2_global + pair_worldid_out[idx] = raw_pair_worldid_in[tid] + + +@wp.kernel +def _flex_selfcollision_narrowphase( + # Model: + nflex: int, + opt_ccd_tolerance: wp.array[float], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elem: wp.array[int], + flex_radius: wp.array[float], + flex_elemflexid: wp.array[int], + # Data in: + flexvert_xpos_in: wp.array2d[wp.vec3], + # In: + max_candidates: int, + gjk_iterations: int, + epa_iterations: int, + npairs_in: wp.array[int], + pair_elem1_in: wp.array[int], + pair_elem2_in: wp.array[int], + pair_worldid_in: wp.array[int], + max_pairs: int, + n_total_elems: int, + # Data out: + overflow_out: wp.array[int], + # Out: + workspace_verts_out: wp.array[wp.vec3], + epa_vert_out: wp.array2d[wp.vec3], + epa_vert_index_out: wp.array2d[int], + epa_face_out: wp.array2d[int], + epa_pr_out: wp.array2d[wp.vec3], + epa_norm2_out: wp.array2d[float], + epa_horizon_out: wp.array2d[int], + cand_dist_out: wp.array[float], + cand_pos_out: wp.array[wp.vec3], + cand_nrm_out: wp.array[wp.vec3], + cand_geom_out: wp.array[wp.vec2i], + cand_flex_out: wp.array[wp.vec2i], + cand_elem_out: wp.array[wp.vec2i], + cand_vert_out: wp.array[wp.vec2i], + cand_worldid_out: wp.array[int], + cand_type_out: wp.array[int], + cand_geomcollisionid_out: wp.array[int], + ncand_out: wp.array[int], +): + """Process SAP-identified pairs through narrowphase (GJK/EPA).""" + pairid = wp.tid() + + # Check bounds + actual_npairs = npairs_in[0] + if pairid >= actual_npairs or pairid >= max_pairs: + return + + elem1_global = pair_elem1_in[pairid] + elem2_global = pair_elem2_in[pairid] + worldid = pair_worldid_in[pairid] + + flexid = flex_elemflexid[elem1_global] + radius = flex_radius[flexid] + dim = flex_dim[flexid] + vert_adr = flex_vertadr[flexid] + elem_adr = flex_elemadr[flexid] + + e1 = elem1_global - elem_adr + e2 = elem2_global - elem_adr + elem_data_idx1 = flex_elemdataadr[flexid] + e1 * (dim + 1) + elem_data_idx2 = flex_elemdataadr[flexid] + e2 * (dim + 1) + + v1_indices = _get_element_vertices(flex_elem, dim, elem_data_idx1) + v2_indices = _get_element_vertices(flex_elem, dim, elem_data_idx2) + + # Workspace for this pair + offset1 = pairid * 8 + for idx in range(dim + 1): + workspace_verts_out[offset1 + idx] = flexvert_xpos_in[worldid, vert_adr + v1_indices[idx]] + + if dim == 1: + # Capsule-capsule collision + p0 = workspace_verts_out[offset1] + p1 = workspace_verts_out[offset1 + 1] + cap1_pos = 0.5 * (p0 + p1) + cap1_axis = wp.normalize(p1 - p0) + cap1_half_len = 0.5 * wp.length(p1 - p0) + + p2_0 = flexvert_xpos_in[worldid, vert_adr + v2_indices[0]] + p2_1 = flexvert_xpos_in[worldid, vert_adr + v2_indices[1]] + cap2_pos = 0.5 * (p2_0 + p2_1) + cap2_axis = wp.normalize(p2_1 - p2_0) + cap2_half_len = 0.5 * wp.length(p2_1 - p2_0) + + margin = 0.0 + + contact_dist, contact_pos, contact_normal = collision_primitive_core.capsule_capsule( + cap1_pos, cap1_axis, radius, cap1_half_len, cap2_pos, cap2_axis, radius, cap2_half_len, margin + ) + + for c in range(2): + d_val = contact_dist[c] + if d_val < 0.0: + _write_candidate_contact( + max_candidates, + d_val, + contact_pos[c], + contact_normal[c], + -2, + flexid, + e1, + e2, + worldid, + overflow_out, + cand_dist_out, + cand_pos_out, + cand_nrm_out, + cand_geom_out, + cand_flex_out, + cand_elem_out, + cand_vert_out, + cand_worldid_out, + cand_type_out, + cand_geomcollisionid_out, + ncand_out, + ) + else: + # GJK/EPA for dim >= 2 + offset2 = pairid * 8 + 4 + for idx in range(dim + 1): + workspace_verts_out[offset2 + idx] = flexvert_xpos_in[worldid, vert_adr + v2_indices[idx]] + + geom1 = Geom() + geom1.pos = wp.vec3(0.0) + geom1.rot = wp.identity(n=3, dtype=float) + geom1.size = wp.vec3(0.0) + geom1.margin = 2.0 * radius + geom1.vert = workspace_verts_out + geom1.vertadr = offset1 + geom1.vertnum = dim + 1 + geom1.graphadr = -1 + geom1.index = -1 + + geom2 = Geom() + geom2.pos = wp.vec3(0.0) + geom2.rot = wp.identity(n=3, dtype=float) + geom2.size = wp.vec3(0.0) + geom2.margin = 2.0 * radius + geom2.vert = workspace_verts_out + geom2.vertadr = offset2 + geom2.vertnum = dim + 1 + geom2.graphadr = -1 + geom2.index = -1 + + center1 = wp.vec3(0.0) + for idx in range(dim + 1): + center1 += workspace_verts_out[offset1 + idx] + center1 = center1 / float(dim + 1) + + center2 = wp.vec3(0.0) + for idx in range(dim + 1): + center2 += workspace_verts_out[offset2 + idx] + center2 = center2 / float(dim + 1) + + tol = opt_ccd_tolerance[0 % opt_ccd_tolerance.shape[0]] + + dist, ncontact, w1, w2, _ = ccd( + tol, + 2.0 * radius, + gjk_iterations, + epa_iterations, + geom1, + geom2, + int(GeomType.MESH), + int(GeomType.MESH), + center1, + center2, + epa_vert_out[pairid], + epa_vert_index_out[pairid], + epa_face_out[pairid], + epa_pr_out[pairid], + epa_norm2_out[pairid], + epa_horizon_out[pairid], + ) + + phys_dist = dist + if ncontact > 0 and phys_dist < 0.0: + p1_0 = workspace_verts_out[offset1] + p1_1 = workspace_verts_out[offset1 + 1] + p1_2 = workspace_verts_out[offset1 + 2] + p2_0 = workspace_verts_out[offset2] + p2_1 = workspace_verts_out[offset2 + 1] + p2_2 = workspace_verts_out[offset2 + 2] + if not (_inside_triangle(w1, p1_0, p1_1, p1_2, 0.2) and _inside_triangle(w2, p2_0, p2_1, p2_2, 0.2)): + return + + pos = 0.5 * (w1 + w2) + nrm = wp.normalize(w1 - w2) + _write_candidate_contact( + max_candidates, + phys_dist, + pos, + nrm, + -2, + flexid, + e1, + e2, + worldid, + overflow_out, + cand_dist_out, + cand_pos_out, + cand_nrm_out, + cand_geom_out, + cand_flex_out, + cand_elem_out, + cand_vert_out, + cand_worldid_out, + cand_type_out, + cand_geomcollisionid_out, + ncand_out, + ) + + @wp.kernel(module="unique", enable_backward=False) def _flex_active_element_collisions_detect( # Model: @@ -1801,7 +2185,7 @@ def _flex_active_element_collisions_detect( geom1 = Geom() geom1.pos = wp.vec3(0.0) - geom1.rot = wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0) + geom1.rot = wp.identity(n=3, dtype=float) geom1.size = wp.vec3(0.0) geom1.margin = 2.0 * radius geom1.vert = workspace_verts_out @@ -1812,7 +2196,7 @@ def _flex_active_element_collisions_detect( geom2 = Geom() geom2.pos = wp.vec3(0.0) - geom2.rot = wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0) + geom2.rot = wp.identity(n=3, dtype=float) geom2.size = wp.vec3(0.0) geom2.margin = 2.0 * radius geom2.vert = workspace_verts_out @@ -1985,6 +2369,7 @@ def _flex_narrowphase_unified( and gtype != int(GeomType.BOX) and gtype != int(GeomType.CYLINDER) and gtype != int(GeomType.MESH) + and gtype != int(GeomType.ELLIPSOID) ): return @@ -2073,6 +2458,103 @@ def _flex_narrowphase_unified( cand_geomcollisionid_out, ncand_out, ) + elif gtype == int(GeomType.ELLIPSOID): + ccdid = wp.atomic_add(nccd, 0, 1) + if ccdid >= naccdmax_in: + if opt_warn_overflow: + wp.printf("CCD overflow in flex narrowphase - please increase naccdmax to %u\n", ccdid) + wp.atomic_or(overflow_out, worldid, wp.static(OverflowType.CCD)) + else: + tolerance = opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]] + + # Construct Ellipsoid Geom (geom1) + geom1 = Geom() + geom1.pos = geom_pos + geom1.rot = geom_rot + geom1.size = geom_size_val + geom1.margin = 0.0 + geom1.index = -1 + + # Construct Triangle Geom (geom2) + geom2 = Geom() + geom2.pos = wp.vec3(0.0, 0.0, 0.0) + geom2.rot = wp.mat33(t1[0], t1[1], t1[2], t2[0], t2[1], t2[2], t3[0], t3[1], t3[2]) + geom2.margin = 0.0 + geom2.index = -1 + + centroid = (t1 + t2 + t3) * wp.static(1.0 / 3.0) + r_geom = wp.length(geom_size_val) + d1 = wp.length(t1 - centroid) + d2 = wp.length(t2 - centroid) + d3 = wp.length(t3 - centroid) + r_tri = wp.max(d1, wp.max(d2, d3)) + + if wp.length(centroid - geom_pos) <= r_geom + r_tri + margin + tri_radius + 0.04: + dist, ncontact, w1, w2, idx = ccd( + tolerance, + margin + tri_radius, + ccd_iterations, + ccd_iterations, + geom1, + geom2, + int(GeomType.ELLIPSOID), + int(GeomType.TRIANGLE), + geom_pos, + centroid, + epa_vert[ccdid], + epa_vert_index[ccdid], + epa_face[ccdid], + epa_pr[ccdid], + epa_norm2[ccdid], + epa_horizon[ccdid], + ) + + if ncontact > 0 and dist < margin + tri_radius: + if _inside_triangle(w2, t1, t2, t3, 0.2): + if dist < 0.0: + normal = wp.normalize(w1 - w2) + else: + normal = wp.normalize(w2 - w1) + + # Project triangle vertices onto normal to find deepest penetration + dist_v0 = wp.dot(t1 - w1, normal) - tri_radius + dist_v1 = wp.dot(t2 - w1, normal) - tri_radius + dist_v2 = wp.dot(t3 - w1, normal) - tri_radius + + min_dist = wp.min(dist_v0, wp.min(dist_v1, dist_v2)) + if min_dist < margin: + deepest_vert = v0_local + pos = t1 - normal * (tri_radius + 0.5 * dist_v0) + if dist_v1 < dist_v0 and dist_v1 < dist_v2: + deepest_vert = v1_local + pos = t2 - normal * (tri_radius + 0.5 * dist_v1) + elif dist_v2 < dist_v0 and dist_v2 < dist_v1: + deepest_vert = v2_local + pos = t3 - normal * (tri_radius + 0.5 * dist_v2) + + _write_candidate_contact( + max_candidates, + min_dist, + pos, + normal, + geomid, + flexid, + local_tri_id, + -1, + worldid, + overflow_out, + cand_dist_out, + cand_pos_out, + cand_nrm_out, + cand_geom_out, + cand_flex_out, + cand_elem_out, + cand_vert_out, + cand_worldid_out, + cand_type_out, + cand_geomcollisionid_out, + ncand_out, + ) else: _collide_geom_triangle_detect( max_candidates, @@ -2105,6 +2587,264 @@ def _flex_narrowphase_unified( ) +@wp.kernel +def _flex_narrowphase_tet_detect( + # Model: + ngeom: int, + nflex: int, + geom_type: wp.array[int], + geom_contype: wp.array[int], + geom_conaffinity: wp.array[int], + geom_size: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + flex_contype: wp.array[int], + flex_conaffinity: wp.array[int], + flex_margin: wp.array[float], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elem: wp.array[int], + flex_radius: wp.array[float], + # Data in: + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], + nworld_in: int, + # In: + max_candidates: int, + # Data out: + overflow_out: wp.array[int], + # Out: + cand_dist_out: wp.array[float], + cand_pos_out: wp.array[wp.vec3], + cand_nrm_out: wp.array[wp.vec3], + cand_geom_out: wp.array[wp.vec2i], + cand_flex_out: wp.array[wp.vec2i], + cand_elem_out: wp.array[wp.vec2i], + cand_vert_out: wp.array[wp.vec2i], + cand_worldid_out: wp.array[int], + cand_type_out: wp.array[int], + cand_geomcollisionid_out: wp.array[int], + ncand_out: wp.array[int], +): + worldid, elemid = wp.tid() + + # Find which flex owns this element + flexid = int(-1) + for i in range(nflex): + if flex_dim[i] != 3: + continue + elem_adr = flex_elemadr[i] + elem_num = flex_elemnum[i] + if elemid >= elem_adr and elemid < elem_adr + elem_num: + flexid = i + break + + if flexid < 0: + return + + vert_adr = flex_vertadr[flexid] + tri_radius = flex_radius[flexid] + tri_margin = flex_margin[flexid] + + # Extract 4 tet vertex indices (dim+1 = 4 for dim=3) + local_elemid = elemid - flex_elemadr[flexid] + edata_idx = flex_elemdataadr[flexid] + local_elemid * 4 + v0 = flex_elem[edata_idx] + v1 = flex_elem[edata_idx + 1] + v2 = flex_elem[edata_idx + 2] + v3 = flex_elem[edata_idx + 3] + + # Fetch world-space vertex positions + p0 = flexvert_xpos_in[worldid, vert_adr + v0] + p1 = flexvert_xpos_in[worldid, vert_adr + v1] + p2 = flexvert_xpos_in[worldid, vert_adr + v2] + p3 = flexvert_xpos_in[worldid, vert_adr + v3] + + # TODO: Add a broadphase + for geomid in range(ngeom): + gtype = geom_type[geomid] + if ( + gtype != int(GeomType.SPHERE) + and gtype != int(GeomType.CAPSULE) + and gtype != int(GeomType.BOX) + and gtype != int(GeomType.CYLINDER) + ): + continue + + g_contype = geom_contype[geomid] + g_conaffinity = geom_conaffinity[geomid] + f_contype = flex_contype[flexid] + f_conaffinity = flex_conaffinity[flexid] + if not ((g_contype & f_conaffinity) or (f_contype & g_conaffinity)): + continue + + geom_margin_val = geom_margin[worldid % geom_margin.shape[0], geomid] + margin = geom_margin_val + tri_margin + + geom_pos = geom_xpos_in[worldid, geomid] + geom_rot = geom_xmat_in[worldid, geomid] + geom_size_val = geom_size[worldid % geom_size.shape[0], geomid] + + # Test all 4 triangular faces of the tet against the geom. + # Face k is the triangle opposite vertex k: + # Face 0: (v1, v2, v3) + # Face 1: (v0, v2, v3) + # Face 2: (v0, v1, v3) + # Face 3: (v0, v1, v2) + for face in range(4): + if face == 0: + t1 = p1 + t2 = p2 + t3 = p3 + elif face == 1: + t1 = p0 + t2 = p2 + t3 = p3 + elif face == 2: + t1 = p0 + t2 = p1 + t3 = p3 + else: + t1 = p0 + t2 = p1 + t3 = p2 + + _collide_geom_triangle_detect( + max_candidates, + gtype, + geom_pos, + geom_rot, + geom_size_val, + t1, + t2, + t3, + tri_radius, + margin, + geomid, + flexid, + elemid, + -1, + worldid, + overflow_out, + cand_dist_out, + cand_pos_out, + cand_nrm_out, + cand_geom_out, + cand_flex_out, + cand_elem_out, + cand_vert_out, + cand_worldid_out, + cand_type_out, + cand_geomcollisionid_out, + ncand_out, + ) + + +@wp.kernel +def _compute_filter_key( + # Model: + ngeom: int, + nflex: int, + # In: + ncand: wp.array[int], + cand_geom: wp.array[wp.vec2i], + cand_flex: wp.array[wp.vec2i], + cand_worldid: wp.array[int], + # Out: + key_out: wp.array[int], + val_out: wp.array[int], +): + """Compute sort key for candidate grouping. + + Groups candidates by (worldid, flex_id, geom_id) so that duplicates + are contiguous after sorting. Self-collision contacts (geom_id < 0) + are mapped to a sentinel value (ngeom). + """ + i = wp.tid() + if i >= ncand[0]: + key_out[i] = 2147483647 # INT_MAX: sort unused entries to end + val_out[i] = i + return + + worldid = cand_worldid[i] + flex_id = cand_flex[i][1] + geom_id = cand_geom[i][0] + + # Map self-collision (geom_id < 0) to sentinel + g = geom_id + if g < 0: + g = ngeom + + key_out[i] = worldid * (nflex + 1) * (ngeom + 2) + flex_id * (ngeom + 2) + g + val_out[i] = i + + +@wp.kernel +def _filter_flex_candidates_sorted( + # In: + ncand: wp.array[int], + epsilon: float, + sort_key: wp.array[int], + sort_val: wp.array[int], + cand_dist: wp.array[float], + cand_pos: wp.array[wp.vec3], + # Out: + cand_active_out: wp.array[int], +): + """Filter duplicate candidates using sorted order. + + After sorting by group key, candidates in the same group are contiguous. + Each candidate only compares with neighbors sharing the same key, reducing + complexity from O(n^2) to O(n * k) where k is the average group size. + """ + si = wp.tid() + if si >= ncand[0]: + return + + i = sort_val[si] + my_key = sort_key[si] + pos_i = cand_pos[i] + dist_i = cand_dist[i] + eps2 = epsilon * epsilon + + keep = int(1) + + # Compare with same-key neighbors (backward) + j = si - 1 + while j >= 0: + if sort_key[j] != my_key: + break + oj = sort_val[j] + diff = pos_i - cand_pos[oj] + if wp.dot(diff, diff) < eps2: + dist_j = cand_dist[oj] + if dist_j < dist_i: + keep = 0 + elif dist_j == dist_i and oj < i: + keep = 0 + j -= 1 + + # Compare with same-key neighbors (forward) + j = si + 1 + while j < ncand[0]: + if sort_key[j] != my_key: + break + oj = sort_val[j] + diff = pos_i - cand_pos[oj] + if wp.dot(diff, diff) < eps2: + dist_j = cand_dist[oj] + if dist_j < dist_i: + keep = 0 + elif dist_j == dist_i and oj < i: + keep = 0 + j += 1 + + cand_active_out[i] = keep + + @wp.kernel def _filter_flex_candidates( # In: @@ -2356,7 +3096,7 @@ def flex_collision(m: Model, d: Data, ctx): # EPA workspaces if mesh or self collisions are possible epa_iterations = m.opt.ccd_iterations - if m.nmesh > 0: + if m.nmesh > 0 or m.has_ellipsoid_geom: mesh_epa_vert = wp.empty(shape=(d.naccdmax, 10 + 2 * epa_iterations), dtype=wp.vec3) mesh_epa_vert_index = wp.empty(shape=(d.naccdmax, 10 + 2 * epa_iterations), dtype=int) mesh_epa_face = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) @@ -2377,6 +3117,50 @@ def flex_collision(m: Model, d: Data, ctx): ncollision_dim3 = wp.zeros(1, dtype=int) ncollision_plane = wp.zeros(1, dtype=int) + if m.has_3d_flex: + wp.launch( + _flex_narrowphase_tet_detect, + dim=(d.nworld, m.nflexelem), + inputs=[ + m.ngeom, + m.nflex, + m.geom_type, + m.geom_contype, + m.geom_conaffinity, + m.geom_size, + m.geom_margin, + m.flex_contype, + m.flex_conaffinity, + m.flex_margin, + m.flex_dim, + m.flex_vertadr, + m.flex_elemadr, + m.flex_elemnum, + m.flex_elemdataadr, + m.flex_elem, + m.flex_radius, + d.geom_xpos, + d.geom_xmat, + d.flexvert_xpos, + d.nworld, + d.naconmax, + ], + outputs=[ + d.overflow, + cand_dist, + cand_pos, + cand_nrm, + cand_geom, + cand_flex, + cand_elem, + cand_vert, + cand_worldid, + cand_type, + cand_geomcollisionid, + ncand, + ], + ) + # Update dynamic flex object bounding boxes flex_broadphase(m, d) @@ -2499,125 +3283,6 @@ def flex_collision(m: Model, d: Data, ctx): ], ) - # 3D Flex Element Collisions - if m.flexshell_geom_pair_filtered.shape[0] > 0: - wp.launch( - _flex_broadphase_unified, - dim=(d.nworld, m.flexshell_geom_pair_filtered.shape[0]), - inputs=[ - m.ngeom, - m.nflex, - m.opt.warn_overflow, - m.geom_type, - m.geom_size, - m.geom_aabb, - m.geom_rbound, - m.geom_margin, - m.flex_margin, - m.flex_dim, - m.flex_vertadr, - m.flex_radius, - d.geom_xpos, - d.geom_xmat, - d.flexvert_xpos, - d.naconmax, - d.flex_aabb_min, - d.flex_aabb_max, - m.flex_shelladr, - m.flex_shelldataadr, - m.flex_shell, - m.flexshell_geom_pair_filtered, - m.flex_shellflexid, - ], - outputs=[ - ncollision_dim3, - d.overflow, - ctx.collision_pair, - ctx.collision_worldid, - ], - ) - - wp.launch( - _flex_narrowphase_unified, - dim=d.naconmax, - inputs=[ - m.ngeom, - m.nflex, - m.opt.ccd_tolerance, - m.opt.warn_overflow, - m.geom_type, - m.geom_condim, - m.geom_dataid, - m.geom_priority, - m.geom_solmix, - m.geom_solref, - m.geom_solimp, - m.geom_size, - m.geom_friction, - m.geom_margin, - m.geom_gap, - m.flex_condim, - m.flex_priority, - m.flex_solmix, - m.flex_solref, - m.flex_solimp, - m.flex_friction, - m.flex_margin, - m.flex_gap, - m.flex_dim, - m.flex_vertadr, - m.flex_radius, - m.mesh_vertadr, - m.mesh_vertnum, - m.mesh_graphadr, - m.mesh_vert, - m.mesh_graph, - m.mesh_pos, - m.mesh_polynormal, - m.mesh_polyvertadr, - m.mesh_polyvert, - m.mesh_polymapadr, - m.mesh_polymapnum, - m.mesh_polymap, - d.geom_xpos, - d.geom_xmat, - d.flexvert_xpos, - d.nworld, - d.naconmax, - d.naccdmax, - ncollision_dim3, - m.flex_shelladr, - m.flex_shelldataadr, - m.flex_shell, - m.flex_shellflexid, - ctx.collision_pair, - ctx.collision_worldid, - mesh_epa_vert, - mesh_epa_vert_index, - mesh_epa_face, - mesh_epa_pr, - mesh_epa_norm2, - mesh_epa_horizon, - mesh_nccd, - epa_iterations, - d.naconmax, - ], - outputs=[ - d.overflow, - cand_dist, - cand_pos, - cand_nrm, - cand_geom, - cand_flex, - cand_elem, - cand_vert, - cand_worldid, - cand_type, - cand_geomcollisionid, - ncand, - ], - ) - # Plane Vertex Collisions if m.flexvert_geom_pair_filtered.shape[0] > 0: wp.launch( @@ -2695,6 +3360,7 @@ def flex_collision(m: Model, d: Data, ctx): d.contact.dim, d.contact.geom, d.contact.flex, + d.contact.elem, d.contact.vert, d.contact.worldid, d.contact.type, @@ -2817,84 +3483,313 @@ def flex_collision(m: Model, d: Data, ctx): selfcollide_enabled = m.has_flex_selfcollide if selfcollide_enabled and m.nflexelem > 0: - workspace_verts = wp.empty(d.nworld * m.nflexelem * 8, dtype=wp.vec3) - epa_iterations = m.opt.ccd_iterations - if m.max_flex_dim > 1: - epa_vert = wp.empty(shape=(d.nworld * m.nflexelem, 10 + 2 * epa_iterations), dtype=wp.vec3) - epa_vert_index = wp.empty(shape=(d.nworld * m.nflexelem, 10 + 2 * epa_iterations), dtype=int) - epa_face = wp.empty(shape=(d.nworld * m.nflexelem, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) - epa_pr = wp.empty(shape=(d.nworld * m.nflexelem, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3) - epa_norm2 = wp.empty(shape=(d.nworld * m.nflexelem, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float) - epa_horizon = wp.empty(shape=(d.nworld * m.nflexelem, MJ_MAX_EPAHORIZON), dtype=int) + + # TODO(team): investigate optimal nflexelem threshold for SAP vs brute-force + if m.nflexelem > 32: + # --- SAP broadphase path --- + nelem = m.nflexelem + nworldelem = d.nworld * nelem + + # Fixed projection direction (same as sap_broadphase in collision_driver.py) + # TODO(team): compute optimal SAP direction + direction = wp.vec3(0.5935, 0.7790, 0.1235) + direction = wp.normalize(direction) + + # Allocate SAP arrays + sap_lower = wp.empty((d.nworld, nelem, 2), dtype=float) + sap_upper = wp.empty((d.nworld, nelem), dtype=float) + sap_sort_index = wp.empty((d.nworld, nelem, 2), dtype=int) + sap_range_arr = wp.empty((d.nworld, nelem), dtype=int) + sap_cumsum = wp.empty((d.nworld, nelem), dtype=int) + sap_seg_index = wp.empty(d.nworld + 1, dtype=int) + elem_aabb_lower = wp.empty((d.nworld, nelem), dtype=wp.vec3) + elem_aabb_upper = wp.empty((d.nworld, nelem), dtype=wp.vec3) + + # Step 1: Project element AABBs onto direction + wp.launch( + _flex_sap_project, + dim=(d.nworld, nelem), + inputs=[ + m.nflex, + m.flex_selfcollide, + m.flex_dim, + m.flex_vertadr, + m.flex_elemadr, + m.flex_elemdataadr, + m.flex_elem, + m.flex_radius, + m.flex_elemflexid, + d.flexvert_xpos, + d.nworld, + nelem, + direction, + ], + outputs=[ + sap_lower.reshape((-1, nelem)), + sap_upper, + sap_sort_index.reshape((-1, nelem)), + elem_aabb_lower, + elem_aabb_upper, + sap_seg_index, + ], + ) + + # Step 2: Sort + wp.utils.segmented_sort_pairs( + sap_lower.reshape((-1, nelem)), + sap_sort_index.reshape((-1, nelem)), + nworldelem, + sap_seg_index, + ) + + # Step 3: Range + wp.launch( + sap_range, + dim=(d.nworld, nelem), + inputs=[ + nelem, + sap_lower.reshape((-1, nelem)), + sap_upper, + sap_sort_index.reshape((-1, nelem)), + ], + outputs=[ + sap_range_arr, + ], + ) + + # Step 4: Prefix sum for load balancing + wp.utils.array_scan( + sap_range_arr.reshape(-1), + sap_cumsum.reshape(-1), + True, + ) + + # Step 5: SAP sweep - output pairs only (no narrowphase) + nsweep = 5 * nworldelem + + npairs = wp.zeros(1, dtype=int) + pair_elem1 = wp.empty(d.naconmax, dtype=int) + pair_elem2 = wp.empty(d.naconmax, dtype=int) + pair_worldid = wp.empty(d.naconmax, dtype=int) + + # Step 5a: Generic SAP sweep (shared with geom broadphase) + raw_npairs = wp.zeros(1, dtype=int) + raw_pair_elem1 = wp.empty(d.naconmax, dtype=int) + raw_pair_elem2 = wp.empty(d.naconmax, dtype=int) + raw_pair_worldid = wp.empty(d.naconmax, dtype=int) + + wp.launch( + sap_sweep, + dim=nsweep, + inputs=[ + nelem, + sap_sort_index.reshape((-1, nelem)), + sap_cumsum.reshape(-1), + nsweep, + elem_aabb_lower, + elem_aabb_upper, + d.naconmax, + ], + outputs=[ + raw_npairs, + raw_pair_elem1, + raw_pair_elem2, + raw_pair_worldid, + ], + ) + + # Step 5b: Filter pairs (flex-specific: selfcollide, shared vertices) + wp.launch( + _flex_sap_filter, + dim=d.naconmax, + inputs=[ + m.flex_selfcollide, + m.flex_dim, + m.flex_vertadr, + m.flex_elemadr, + m.flex_elemnum, + m.flex_elemdataadr, + m.flex_vertbodyid, + m.flex_elem, + m.flex_elemflexid, + raw_npairs, + raw_pair_elem1, + raw_pair_elem2, + raw_pair_worldid, + ], + outputs=[ + npairs, + pair_elem1, + pair_elem2, + pair_worldid, + ], + ) + + # Step 6: Narrowphase on actual pairs only + workspace_verts = wp.empty(d.naconmax * 8, dtype=wp.vec3) + + if m.max_flex_dim > 1: + epa_vert = wp.empty(shape=(d.naconmax, 10 + 2 * epa_iterations), dtype=wp.vec3) + epa_vert_index = wp.empty(shape=(d.naconmax, 10 + 2 * epa_iterations), dtype=int) + epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) + epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3) + epa_norm2 = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float) + epa_horizon = wp.empty(shape=(d.naconmax, MJ_MAX_EPAHORIZON), dtype=int) + else: + epa_vert = wp.empty(shape=(1, 1), dtype=wp.vec3) + epa_vert_index = wp.empty(shape=(1, 1), dtype=int) + epa_face = wp.empty(shape=(1, 1), dtype=int) + epa_pr = wp.empty(shape=(1, 1), dtype=wp.vec3) + epa_norm2 = wp.empty(shape=(1, 1), dtype=float) + epa_horizon = wp.empty(shape=(1, 1), dtype=int) + + wp.launch( + _flex_selfcollision_narrowphase, + dim=d.naconmax, + inputs=[ + m.nflex, + m.opt.ccd_tolerance, + m.flex_dim, + m.flex_vertadr, + m.flex_elemadr, + m.flex_elemdataadr, + m.flex_elem, + m.flex_radius, + m.flex_elemflexid, + d.flexvert_xpos, + d.naconmax, + m.opt.ccd_iterations, + epa_iterations, + npairs, + pair_elem1, + pair_elem2, + pair_worldid, + d.naconmax, + m.nflexelem, + ], + outputs=[ + d.overflow, + workspace_verts, + epa_vert, + epa_vert_index, + epa_face, + epa_pr, + epa_norm2, + epa_horizon, + cand_dist, + cand_pos, + cand_nrm, + cand_geom, + cand_flex, + cand_elem, + cand_vert, + cand_worldid, + cand_type, + cand_geomcollisionid, + ncand, + ], + ) + else: - epa_vert = wp.empty(shape=(1, 1), dtype=wp.vec3) - epa_vert_index = wp.empty(shape=(1, 1), dtype=int) - epa_face = wp.empty(shape=(1, 1), dtype=int) - epa_pr = wp.empty(shape=(1, 1), dtype=wp.vec3) - epa_norm2 = wp.empty(shape=(1, 1), dtype=float) - epa_horizon = wp.empty(shape=(1, 1), dtype=int) + # --- Brute-force fallback for small element counts --- + workspace_verts = wp.empty(d.nworld * m.nflexelem * 8, dtype=wp.vec3) - wp.launch( - _flex_active_element_collisions_detect, - dim=(d.nworld, m.nflexelem), - inputs=[ - m.nflex, - m.opt.ccd_tolerance, - m.flex_selfcollide, - m.flex_dim, - m.flex_vertadr, - m.flex_elemadr, - m.flex_elemnum, - m.flex_elemdataadr, - m.flex_vertbodyid, - m.flex_elem, - m.flex_radius, - m.flex_elemflexid, - d.flexvert_xpos, - d.naconmax, - m.opt.ccd_iterations, - epa_iterations, - m.nflexelem, - ], - outputs=[ - d.overflow, - workspace_verts, - epa_vert, - epa_vert_index, - epa_face, - epa_pr, - epa_norm2, - epa_horizon, - cand_dist, - cand_pos, - cand_nrm, - cand_geom, - cand_flex, - cand_elem, - cand_vert, - cand_worldid, - cand_type, - cand_geomcollisionid, - ncand, - ], - ) + if m.max_flex_dim > 1: + epa_vert = wp.empty(shape=(d.nworld * m.nflexelem, 10 + 2 * epa_iterations), dtype=wp.vec3) + epa_vert_index = wp.empty(shape=(d.nworld * m.nflexelem, 10 + 2 * epa_iterations), dtype=int) + epa_face = wp.empty(shape=(d.nworld * m.nflexelem, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) + epa_pr = wp.empty(shape=(d.nworld * m.nflexelem, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3) + epa_norm2 = wp.empty(shape=(d.nworld * m.nflexelem, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float) + epa_horizon = wp.empty(shape=(d.nworld * m.nflexelem, MJ_MAX_EPAHORIZON), dtype=int) + else: + epa_vert = wp.empty(shape=(1, 1), dtype=wp.vec3) + epa_vert_index = wp.empty(shape=(1, 1), dtype=int) + epa_face = wp.empty(shape=(1, 1), dtype=int) + epa_pr = wp.empty(shape=(1, 1), dtype=wp.vec3) + epa_norm2 = wp.empty(shape=(1, 1), dtype=float) + epa_horizon = wp.empty(shape=(1, 1), dtype=int) - # Filter duplicate contacts (e.g. from shared vertices or edges) - cand_active = wp.empty(d.naconmax, dtype=int) + wp.launch( + _flex_active_element_collisions_detect, + dim=(d.nworld, m.nflexelem), + inputs=[ + m.nflex, + m.opt.ccd_tolerance, + m.flex_selfcollide, + m.flex_dim, + m.flex_vertadr, + m.flex_elemadr, + m.flex_elemnum, + m.flex_elemdataadr, + m.flex_vertbodyid, + m.flex_elem, + m.flex_radius, + m.flex_elemflexid, + d.flexvert_xpos, + d.naconmax, + m.opt.ccd_iterations, + epa_iterations, + m.nflexelem, + ], + outputs=[ + d.overflow, + workspace_verts, + epa_vert, + epa_vert_index, + epa_face, + epa_pr, + epa_norm2, + epa_horizon, + cand_dist, + cand_pos, + cand_nrm, + cand_geom, + cand_flex, + cand_elem, + cand_vert, + cand_worldid, + cand_type, + cand_geomcollisionid, + ncand, + ], + ) + + # Filter duplicate contacts using sort-based deduplication. + # Sort candidates by (worldid, flex_id, geom_id) so duplicates are contiguous, + # then compare only within each group. This is O(n log n) vs the naive O(n^2). + filter_key = wp.empty(d.naconmax * 2, dtype=int) + filter_val = wp.empty(d.naconmax * 2, dtype=int) wp.launch( - _filter_flex_candidates, + _compute_filter_key, dim=d.naconmax, inputs=[ - d.naconmax, + m.ngeom, + m.nflex, ncand, - 1e-3, # epsilon - cand_dist, - cand_pos, cand_geom, cand_flex, cand_worldid, ], + outputs=[ + filter_key, + filter_val, + ], + ) + wp.utils.radix_sort_pairs(filter_key, filter_val, d.naconmax) + + cand_active = wp.empty(d.naconmax, dtype=int) + wp.launch( + _filter_flex_candidates_sorted, + dim=d.naconmax, + inputs=[ + ncand, + 1e-3, # epsilon + filter_key, + filter_val, + cand_dist, + cand_pos, + ], outputs=[ cand_active, ], diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py index ea16b361..7776ec99 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py @@ -78,6 +78,9 @@ class Polytope: vert_index: wp.array[int] nvert: int + # center point of polytope + center: wp.vec3 + # faces in polytope # 10 bits per each vertex index, while the last significant bits are for # invalid and deleted face @@ -229,6 +232,10 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float: if ret: return 0.0 + # ensure projection points outward from the polytope + if wp.dot(r, p1 - pt.center) < 0.0: + r = -r + face = v1 + (v2 << 10) + (v3 << 20) pt.face[idx] = face pt.face_pr[idx] = r @@ -255,17 +262,14 @@ def _epa_support( @wp.func -def _linear_combine(n: int, coefs: wp.vec4, mat: mat43) -> wp.vec3: - v = wp.vec3(0.0) +def _linear_combine(n: int, scl: wp.vec4, mat: mat43) -> wp.vec3: if n == 1: - v = coefs[0] * mat[0] - elif n == 2: - v = coefs[0] * mat[0] + coefs[1] * mat[1] - elif n == 3: - v = coefs[0] * mat[0] + coefs[1] * mat[1] + coefs[2] * mat[2] - else: - v = coefs[0] * mat[0] + coefs[1] * mat[1] + coefs[2] * mat[2] + coefs[3] * mat[3] - return v + return scl[0] * mat[0] + if n == 2: + return scl[0] * mat[0] + scl[1] * mat[1] + if n == 3: + return scl[0] * mat[0] + scl[1] * mat[1] + scl[2] * mat[2] + return scl[0] * mat[0] + scl[1] * mat[1] + scl[2] * mat[2] + scl[3] * mat[3] @wp.func @@ -278,11 +282,11 @@ def _subdistance(n: int, simplex: mat43) -> wp.vec4: if n == 4: return _S3D(simplex[0], simplex[1], simplex[2], simplex[3]) if n == 3: - coordinates3 = _S2D(simplex[0], simplex[1], simplex[2]) - return wp.vec4(coordinates3[0], coordinates3[1], coordinates3[2], 0.0) + lmbda3 = _S2D(simplex[0], simplex[1], simplex[2]) + return wp.vec4(lmbda3[0], lmbda3[1], lmbda3[2], 0.0) if n == 2: - coordinates2 = _S1D(simplex[0], simplex[1]) - return wp.vec4(coordinates2[0], coordinates2[1], 0.0, 0.0) + lmbda2 = _S1D(simplex[0], simplex[1]) + return wp.vec4(lmbda2[0], lmbda2[1], 0.0, 0.0) return wp.vec4(1.0, 0.0, 0.0, 0.0) @@ -368,51 +372,51 @@ def _S3D(s1: wp.vec3, s2: wp.vec3, s3: wp.vec3, s4: wp.vec3) -> wp.vec4: return wp.vec4(C41 / m_det, C42 / m_det, C43 / m_det, C44 / m_det) # find the smallest distance, and use the corresponding barycentric coordinates - coordinates = wp.vec4(0.0, 0.0, 0.0, 0.0) + lmbda = wp.vec4(0.0, 0.0, 0.0, 0.0) dmin = FLOAT_MAX if not comp1: - subcoord = _S2D(s2, s3, s4) - x = subcoord[0] * s2 + subcoord[1] * s3 + subcoord[2] * s4 + sublmbda = _S2D(s2, s3, s4) + x = sublmbda[0] * s2 + sublmbda[1] * s3 + sublmbda[2] * s4 d = wp.dot(x, x) - coordinates[0] = 0.0 - coordinates[1] = subcoord[0] - coordinates[2] = subcoord[1] - coordinates[3] = subcoord[2] + lmbda[0] = 0.0 + lmbda[1] = sublmbda[0] + lmbda[2] = sublmbda[1] + lmbda[3] = sublmbda[2] dmin = d if not comp2: - subcoord = _S2D(s1, s3, s4) - x = subcoord[0] * s1 + subcoord[1] * s3 + subcoord[2] * s4 + sublmbda = _S2D(s1, s3, s4) + x = sublmbda[0] * s1 + sublmbda[1] * s3 + sublmbda[2] * s4 d = wp.dot(x, x) if d < dmin: - coordinates[0] = subcoord[0] - coordinates[1] = 0.0 - coordinates[2] = subcoord[1] - coordinates[3] = subcoord[2] + lmbda[0] = sublmbda[0] + lmbda[1] = 0.0 + lmbda[2] = sublmbda[1] + lmbda[3] = sublmbda[2] dmin = d if not comp3: - subcoord = _S2D(s1, s2, s4) - x = subcoord[0] * s1 + subcoord[1] * s2 + subcoord[2] * s4 + sublmbda = _S2D(s1, s2, s4) + x = sublmbda[0] * s1 + sublmbda[1] * s2 + sublmbda[2] * s4 d = wp.dot(x, x) if d < dmin: - coordinates[0] = subcoord[0] - coordinates[1] = subcoord[1] - coordinates[2] = 0.0 - coordinates[3] = subcoord[2] + lmbda[0] = sublmbda[0] + lmbda[1] = sublmbda[1] + lmbda[2] = 0.0 + lmbda[3] = sublmbda[2] dmin = d if not comp4: - subcoord = _S2D(s1, s2, s3) - x = subcoord[0] * s1 + subcoord[1] * s2 + subcoord[2] * s3 + sublmbda = _S2D(s1, s2, s3) + x = sublmbda[0] * s1 + sublmbda[1] * s2 + sublmbda[2] * s3 d = wp.dot(x, x) if d < dmin: - coordinates[0] = subcoord[0] - coordinates[1] = subcoord[1] - coordinates[2] = subcoord[2] - coordinates[3] = 0.0 - return coordinates + lmbda[0] = sublmbda[0] + lmbda[1] = sublmbda[1] + lmbda[2] = sublmbda[2] + lmbda[3] = 0.0 + return lmbda @wp.func @@ -528,36 +532,36 @@ def _S2D(s1: wp.vec3, s2: wp.vec3, s3: wp.vec3) -> wp.vec3: # find the smallest distance, and use the corresponding barycentric coordinates dmin = FLOAT_MAX - coordinates = wp.vec3(0.0, 0.0, 0.0) + lmbda = wp.vec3(0.0, 0.0, 0.0) if not comp1: - subcoord = _S1D(s2, s3) - x = subcoord[0] * s2 + subcoord[1] * s3 + sublmbda = _S1D(s2, s3) + x = sublmbda[0] * s2 + sublmbda[1] * s3 d = wp.dot(x, x) - coordinates[0] = 0.0 - coordinates[1] = subcoord[0] - coordinates[2] = subcoord[1] + lmbda[0] = 0.0 + lmbda[1] = sublmbda[0] + lmbda[2] = sublmbda[1] dmin = d if not comp2: - subcoord = _S1D(s1, s3) - x = subcoord[0] * s1 + subcoord[1] * s3 + sublmbda = _S1D(s1, s3) + x = sublmbda[0] * s1 + sublmbda[1] * s3 d = wp.dot(x, x) if d < dmin: - coordinates[0] = subcoord[0] - coordinates[1] = 0.0 - coordinates[2] = subcoord[1] + lmbda[0] = sublmbda[0] + lmbda[1] = 0.0 + lmbda[2] = sublmbda[1] dmin = d if not comp3: - subcoord = _S1D(s1, s2) - x = subcoord[0] * s1 + subcoord[1] * s2 + sublmbda = _S1D(s1, s2) + x = sublmbda[0] * s1 + sublmbda[1] * s2 d = wp.dot(x, x) if d < dmin: - coordinates[0] = subcoord[0] - coordinates[1] = subcoord[1] - coordinates[2] = 0.0 - return coordinates + lmbda[0] = sublmbda[0] + lmbda[1] = sublmbda[1] + lmbda[2] = 0.0 + return lmbda @wp.func @@ -566,13 +570,18 @@ def _S1D(s1: wp.vec3, s2: wp.vec3) -> wp.vec2: p_o = _project_origin_line(s1, s2) # find the axis with the largest projection "shadow" of the simplex - mu_max = 0.0 + mu_max = s1[0] - s2[0] index = 0 - for i in range(3): - mu = s1[i] - s2[i] - if wp.abs(mu) >= wp.abs(mu_max): - mu_max = mu - index = i + + mu = s1[1] - s2[1] + if wp.abs(mu) >= wp.abs(mu_max): + mu_max = mu + index = 1 + + mu = s1[2] - s2[2] + if wp.abs(mu) >= wp.abs(mu_max): + mu_max = mu + index = 2 C1 = p_o[index] - s2[index] C2 = s1[index] - p_o[index] @@ -605,21 +614,21 @@ def gjk( simplex_index1 = wp.vec4i() simplex_index2 = wp.vec4i() n = int(0) - coordinates = wp.vec4() # barycentric coordinates + lmbda = wp.vec4() # barycentric coordinates tol2 = tolerance * tolerance epsilon = wp.where(is_discrete, 0.0, 0.5 * tol2) # set initial guess x_k = x1_0 - x2_0 - xnorm_old = FLOAT_MAX + xnorm2_old = FLOAT_MAX for _ in range(gjk_iterations): - xnorm = wp.dot(x_k, x_k) + xnorm2 = wp.dot(x_k, x_k) # TODO(kbayes): determine new constant here - if xnorm < tol2 or wp.abs(xnorm_old - xnorm) < tol2: + if xnorm2 < tol2 or wp.abs(xnorm2_old - xnorm2) < tol2: break - xnorm_old = xnorm - dir_neg = x_k / wp.sqrt(xnorm) + xnorm2_old = xnorm2 + dir_neg = x_k / wp.sqrt(xnorm2) # compute kth support point in geom1 sp = support(geom1, geomtype1, -dir_neg) @@ -636,6 +645,11 @@ def gjk( # compute the kth support point simplex[n] = simplex1[n] - simplex2[n] + # stopping criteria using the Frank-Wolfe duality gap given by + # |f(x_k) - f(x_min)|^2 <= < grad f(x_k), (x_k - simplex[n]) > + if wp.dot(x_k, x_k - simplex[n]) < epsilon: + break + if cutoff == 0.0: if wp.dot(x_k, simplex[n]) > 0.0: result = GJKResult() @@ -644,25 +658,20 @@ def gjk( return result elif cutoff < FLOAT_MAX: vs = wp.dot(x_k, simplex[n]) - if wp.dot(x_k, simplex[n]) > 0.0 and (vs * vs / xnorm) >= cutoff2: + if wp.dot(x_k, simplex[n]) > 0.0 and (vs * vs / xnorm2) >= cutoff2: result = GJKResult() result.dim = 0 result.dist = FLOAT_MAX return result - # stopping criteria using the Frank-Wolfe duality gap given by - # |f(x_k) - f(x_min)|^2 <= < grad f(x_k), (x_k - simplex[n]) > - if wp.dot(x_k, x_k - simplex[n]) < epsilon: - break - # run the distance subalgorithm to compute the barycentric coordinates # of the closest point to the origin in the simplex - coordinates = _subdistance(n + 1, simplex) + lmbda = _subdistance(n + 1, simplex) # remove vertices from the simplex no longer needed n = int(0) for i in range(4): - if coordinates[i] == 0.0: + if lmbda[i] == 0.0: continue simplex[n] = simplex[i] @@ -670,7 +679,7 @@ def gjk( simplex2[n] = simplex2[i] simplex_index1[n] = simplex_index1[i] simplex_index2[n] = simplex_index2[i] - coordinates[n] = coordinates[i] + lmbda[n] = lmbda[i] n += int(1) # SHOULD NOT OCCUR @@ -678,7 +687,7 @@ def gjk( break # get the next iteration of x_k - x_next = _linear_combine(n, coordinates, simplex) + x_next = _linear_combine(n, lmbda, simplex) # x_k has converged to minimum if _almost_equal(x_next, x_k): @@ -696,8 +705,8 @@ def gjk( # compute the approximate witness points # if n is zero, then there was an immediate return meaning the initial points # are the witness points - result.x1 = wp.where(n == 0, x1_0, _linear_combine(n, coordinates, simplex1)) - result.x2 = wp.where(n == 0, x2_0, _linear_combine(n, coordinates, simplex2)) + result.x1 = wp.where(n == 0, x1_0, _linear_combine(n, lmbda, simplex1)) + result.x2 = wp.where(n == 0, x2_0, _linear_combine(n, lmbda, simplex2)) result.dist = wp.norm_l2(x_k) result.dim = n @@ -974,6 +983,9 @@ def _polytope2( """Create polytope for EPA given a 1-simplex from GJK.""" diff = simplex[1] - simplex[0] + # set the polytope center + pt.center = 0.5 * (simplex[0] + simplex[1]) + # find component with smallest magnitude (so cross product is largest) value = FLOAT_MAX index = 0 @@ -1063,6 +1075,9 @@ def _polytope3( geomtype2: int, ) -> Polytope: """Create polytope for EPA given a 2-simplex from GJK.""" + # set the polytope center + pt.center = (simplex[0] + simplex[1] + simplex[2]) * wp.static(1.0 / 3.0) + # get normals in both directions n = wp.cross(simplex[1] - simplex[0], simplex[2] - simplex[0]) if wp.norm_l2(n) < MINVAL: @@ -1146,6 +1161,9 @@ def _polytope4( simplex_index2: wp.vec4i, ) -> Tuple[Polytope, GJKResult]: """Create polytope for EPA given a 3-simplex from GJK.""" + # set the polytope center + pt.center = 0.25 * (simplex[0] + simplex[1] + simplex[2] + simplex[3]) + pt.vert[0] = simplex1[0] pt.vert[1] = simplex2[0] pt.vert[2] = simplex1[1] @@ -1269,15 +1287,15 @@ def _epa( # compute support point w from the closest face's normal lower = wp.sqrt(lower2) wi = pt.nvert - face_pr_normalized = pt.face_pr[idx] / lower - i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr_normalized) + face_pr = pt.face_pr[idx] + i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr / lower) w = pt.vert[2 * wi] - pt.vert[2 * wi + 1] geom1.index = i1 geom2.index = i2 pt.nvert += 1 - # upper bound for kth iteration - upper_k = wp.dot(face_pr_normalized, w) + # upper bound for kth iteration (dot product before normalizing for better precision) + upper_k = wp.dot(face_pr, w) / lower if upper_k < upper: upper = upper_k upper2 = upper * upper diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py index 42476838..67f5c401 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py @@ -808,6 +808,7 @@ def _equality_flex(is_sparse: bool, newton: bool): nv: int, opt_timestep: wp.array[float], opt_disableflags: int, + flex_interp: wp.array[int], flex_edgeadr: wp.array[int], flex_edgenum: wp.array[int], flexedge_length0: wp.array[float], @@ -854,6 +855,11 @@ def _equality_flex(is_sparse: bool, newton: bool): return flexid = eq_obj1id[eqid] + + # skip edge-based equality for interpolated (trilinear/quadratic) flex + if flex_interp[flexid] != 0: + return + if edgeid < flex_edgeadr[flexid] or edgeid >= flex_edgeadr[flexid] + flex_edgenum[flexid]: return @@ -1406,6 +1412,273 @@ def _equality_weld(is_sparse: bool, newton: bool): return kernel +@cache_kernel +def _equality_flexstrain(is_sparse: bool, newton: bool): + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # Model: + nv: int, + opt_timestep: wp.array[float], + opt_disableflags: int, + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + body_invweight0: wp.array2d[wp.vec2], + dof_bodyid: wp.array[int], + flex_interp: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], + flex_nodenum: wp.array[int], + flex_stiffnessadr: wp.array[int], + flex_nodebodyid: wp.array[int], + flex_node: wp.array[wp.vec3], + flex_node0: wp.array[wp.vec3], + flex_stiffness: wp.array[float], + flex_centered: wp.array[bool], + eq_obj1id: wp.array[int], + eq_solref: wp.array2d[wp.vec2], + eq_solimp: wp.array2d[vec5], + eq_data: wp.array2d[vec11], + body_isdofancestor: wp.array2d[int], + eq_flexstrain_adr: wp.array[int], + flexstrain_J_rownnz: wp.array[int], + flexstrain_J_rowadr: wp.array[int], + flexstrain_J_colind: wp.array[int], + # Data in: + qvel_in: wp.array2d[float], + eq_active_in: wp.array2d[bool], + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + njmax_in: int, + njmax_nnz_in: int, + flexnode_xpos_in: wp.array2d[wp.vec3], + # Data out: + ne_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_jtdaj_adr_out: wp.array2d[int], + efc_jtdaj_nrow_out: wp.array2d[int], + efc_jtdaj_nblock_out: wp.array[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], + # Out: + efc_nnz_out: wp.array[int], + ): + worldid, eqstrainid = wp.tid() + eqid = eq_flexstrain_adr[eqstrainid] + + if not eq_active_in[worldid, eqid]: + return + + f = eq_obj1id[eqid] + order = flex_interp[f] + if order <= 0: + return + + # nodes per cell + npc = (order + 1) * (order + 1) * (order + 1) + + # cell indices from eq_data + data = eq_data[worldid % eq_data.shape[0], eqid] + ci = int(data[0]) + cj = int(data[1]) + ck = int(data[2]) + + cellnum = flex_cellnum[f] + cy = cellnum[1] + cz = cellnum[2] + nstart = flex_nodeadr[f] + ny_g = cy * order + 1 + nz_g = cz * order + 1 + + ndof_cell = 3 * npc + + # read eigenmode data from flex_stiffness + cell_idx = ci * cy * cz + cj * cz + ck + k_base = flex_stiffnessadr[f] + cell_idx * ndof_cell * ndof_cell + neig = int(flex_stiffness[k_base]) + + if neig == 0: + return + + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + + # gather cell node positions and reference positions (order==1: npc=8) + # We compute the corotational quaternion from the deformation gradient + # at the cell center (0.5, 0.5, 0.5) + + cell_quat = support.compute_interp_cell_quat(flexnode_xpos_in, order, ci, cj, ck, cy, cz, ny_g, nz_g, nstart, worldid) + cell_quat_inv = wp.quat(-cell_quat[0], -cell_quat[1], -cell_quat[2], cell_quat[3]) + + # Compute average invweight across cell nodes (translation component) + avg_invweight = float(0.0) + idx_iw = int(0) + for li_iw in range(order + 1): + for lj_iw in range(order + 1): + for lk_iw in range(order + 1): + if idx_iw < npc: + gi_iw = ci * order + li_iw + gj_iw = cj * order + lj_iw + gk_iw = ck * order + lk_iw + gidx_iw = gi_iw * ny_g * nz_g + gj_iw * nz_g + gk_iw + bodyid_iw = flex_nodebodyid[nstart + gidx_iw] + avg_invweight += body_invweight0[worldid % body_invweight0.shape[0], bodyid_iw][0] + idx_iw += 1 + avg_invweight /= float(npc) + + # Loop over eigenmodes + for eig in range(neig): + wp.atomic_add(ne_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= njmax_in: + return + + if wp.static(is_sparse and newton): + jgid = wp.atomic_add(efc_jtdaj_nblock_out, worldid, 1) + efc_jtdaj_adr_out[worldid, jgid] = efcid + efc_jtdaj_nrow_out[worldid, jgid] = 6 + + # Read eigenvector from stiffness data + eigvec_base = k_base + 1 + eig * ndof_cell + + # Compute constraint residual: dot(eigvec, displacement_in_corot_frame) + residual = float(0.0) + idx2 = int(0) + for li2 in range(order + 1): + for lj2 in range(order + 1): + for lk2 in range(order + 1): + if idx2 < npc: + gi2 = ci * order + li2 + gj2 = cj * order + lj2 + gk2 = ck * order + lk2 + gidx2 = gi2 * ny_g * nz_g + gj2 * nz_g + gk2 + + xpos_n = flexnode_xpos_in[worldid, nstart + gidx2] + refpos_n = flex_node0[nstart + gidx2] + + # Rotate xpos to corotational frame + xrot = wp.quat_rotate(cell_quat_inv, xpos_n) + displ = wp.vec3(xrot[0] - refpos_n[0], xrot[1] - refpos_n[1], xrot[2] - refpos_n[2]) + + for c in range(3): + residual += flex_stiffness[eigvec_base + 3 * idx2 + c] * displ[c] + + idx2 += 1 + + efc_pos_out[worldid, efcid] = residual + + # Compute Jacobian: J_q = dSdx_world * dxpos/dq + # dSdx_world = R_inv * eigvec (rotate eigenvector to world frame) + # Then contract with per-node body Jacobians + + Jqvel = float(0.0) + invweight = float(0.0) + + rownnz = flexstrain_J_rownnz[eqstrainid] + fs_rowadr = flexstrain_J_rowadr[eqstrainid] + + efc_rowadr = int(0) + if wp.static(is_sparse): + efc_J_rownnz_out[worldid, efcid] = rownnz + efc_rowadr = wp.atomic_add(efc_nnz_out, worldid, rownnz) + if efc_rowadr + rownnz > njmax_nnz_in: + return + efc_J_rowadr_out[worldid, efcid] = efc_rowadr + else: + for q in range(nv): + efc_J_out[worldid, efcid, q] = 0.0 + + for sparseid in range(rownnz): + q = flexstrain_J_colind[fs_rowadr + sparseid] + J_val = float(0.0) + idx3 = int(0) + for li3 in range(order + 1): + for lj3 in range(order + 1): + for lk3 in range(order + 1): + if idx3 < npc: + gi3 = ci * order + li3 + gj3 = cj * order + lj3 + gk3 = ck * order + lk3 + gidx3 = gi3 * ny_g * nz_g + gj3 * nz_g + gk3 + + bodyid3 = flex_nodebodyid[nstart + gidx3] + xpos_n3 = flexnode_xpos_in[worldid, nstart + gidx3] + + if body_isdofancestor[bodyid3, q] != 0: + # Rotate eigenvector component to world frame + eig_local = wp.vec3( + flex_stiffness[eigvec_base + 3 * idx3 + 0], + flex_stiffness[eigvec_base + 3 * idx3 + 1], + flex_stiffness[eigvec_base + 3 * idx3 + 2], + ) + eig_world = wp.quat_rotate(cell_quat, eig_local) + + jacp, jacr = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + body_isdofancestor, + subtree_com_in, + cdof_in, + xpos_n3, + bodyid3, + q, + worldid, + ) + J_val += wp.dot(eig_world, jacp) + idx3 += 1 + + if wp.static(is_sparse): + efc_J_colind_out[worldid, 0, efc_rowadr + sparseid] = q + efc_J_out[worldid, 0, efc_rowadr + sparseid] = J_val + else: + efc_J_out[worldid, efcid, q] = J_val + Jqvel += J_val * qvel_in[worldid, q] + + invweight = avg_invweight + + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + residual, + residual, + invweight, + solref, + solimp, + 0.0, + Jqvel, + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) + + return kernel + + @cache_kernel def _friction_dof(is_sparse: bool, newton: bool): @wp.kernel(module="unique", enable_backward=False) @@ -2021,12 +2294,16 @@ def _get_contact_bodies_and_weights( # Model: geom_bodyid: wp.array[int], flex_dim: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], flex_vertadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_shelldataadr: wp.array[int], + flex_nodebodyid: wp.array[int], flex_vertbodyid: wp.array[int], flex_elem: wp.array[int], flex_shell: wp.array[int], + flex_vert0: wp.array[wp.vec3], # Data in: flexvert_xpos_in: wp.array2d[wp.vec3], # In: @@ -2052,7 +2329,70 @@ def _get_contact_bodies_and_weights( flex_vert_start = flex_vertadr[flex_id] if vert_id >= 0: body = flex_vertbodyid[flex_vert_start + vert_id] - return wp.vec4i(body, -1, -1, -1), wp.vec4(1.0, 0.0, 0.0, 0.0) + if body >= 0: + return wp.vec4i(body, -1, -1, -1), wp.vec4(1.0, 0.0, 0.0, 0.0) + + # Interpolated flex vertex contact: interpolate to the 8 corner nodes + v_adr = flex_vert_start + vert_id + coord = flex_vert0[v_adr] + cn = flex_cellnum[flex_id] + cx = cn[0] + cy = cn[1] + cz = cn[2] + + ci = wp.min(int(coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[flex_id] + + body_ids = wp.vec4i(-1, -1, -1, -1) + weights = wp.vec4(0.0, 0.0, 0.0, 0.0) + n_collected = int(0) + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5 and n_collected < 4: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = flex_nodebodyid[nstart + node_idx] + + # Put into body_ids and weights + if n_collected == 0: + body_ids = wp.vec4i(b, -1, -1, -1) + weights = wp.vec4(w, 0.0, 0.0, 0.0) + elif n_collected == 1: + body_ids = wp.vec4i(body_ids[0], b, -1, -1) + weights = wp.vec4(weights[0], w, 0.0, 0.0) + elif n_collected == 2: + body_ids = wp.vec4i(body_ids[0], body_ids[1], b, -1) + weights = wp.vec4(weights[0], weights[1], w, 0.0) + else: + body_ids = wp.vec4i(body_ids[0], body_ids[1], body_ids[2], b) + weights = wp.vec4(weights[0], weights[1], weights[2], w) + + n_collected += 1 + + # Normalize weights so they sum to 1.0 + w_sum = weights[0] + weights[1] + weights[2] + weights[3] + if w_sum > 1.0e-5: + weights = wp.vec4(weights[0] / w_sum, weights[1] / w_sum, weights[2] / w_sum, weights[3] / w_sum) + + return body_ids, weights # Element contact: Retrieve local vertices dim = flex_dim[flex_id] @@ -2115,16 +2455,106 @@ def _get_contact_bodies_and_weights( w3 = w3 / w_sum b0 = flex_vertbodyid[flex_vert_start + v0] - b1 = flex_vertbodyid[flex_vert_start + v1] - b2 = flex_vertbodyid[flex_vert_start + v2] - b3 = flex_vertbodyid[flex_vert_start + v3] + if b0 >= 0: + b1 = flex_vertbodyid[flex_vert_start + v1] + b2 = flex_vertbodyid[flex_vert_start + v2] + b3 = flex_vertbodyid[flex_vert_start + v3] + return wp.vec4i(b0, b1, b2, b3), wp.vec4(w0, w1, w2, w3) + else: + # Interpolated 3D flex element contact: accumulate node weights and select top 4 + coord0 = flex_vert0[flex_vert_start + v0] + cn = flex_cellnum[flex_id] + cx = cn[0] + cy = cn[1] + cz = cn[2] - return wp.vec4i(b0, b1, b2, b3), wp.vec4(w0, w1, w2, w3) + ci = wp.min(int(coord0[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord0[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord0[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[flex_id] + + # 8 accumulated weights for the voxel corners + W0 = float(0.0) + W1 = float(0.0) + W2 = float(0.0) + W3 = float(0.0) + W4 = float(0.0) + W5 = float(0.0) + W6 = float(0.0) + W7 = float(0.0) + + # Loop over the 4 vertices of the tetrahedron + for k in range(4): + v_idx = v0 + w_k = w0 + if k == 1: + v_idx = v1 + w_k = w1 + elif k == 2: + v_idx = v2 + w_k = w2 + elif k == 3: + v_idx = v3 + w_k = w3 + + coord = flex_vert0[flex_vert_start + v_idx] + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + # Accumulate trilinear weights + W0 += w_k * support.eval_basis_trilinear(local, 0) + W1 += w_k * support.eval_basis_trilinear(local, 1) + W2 += w_k * support.eval_basis_trilinear(local, 2) + W3 += w_k * support.eval_basis_trilinear(local, 3) + W4 += w_k * support.eval_basis_trilinear(local, 4) + W5 += w_k * support.eval_basis_trilinear(local, 5) + W6 += w_k * support.eval_basis_trilinear(local, 6) + W7 += w_k * support.eval_basis_trilinear(local, 7) + + # 8 node body IDs corresponding to the 8 voxel corners + b0_node = flex_nodebodyid[nstart + (ci + 0) * ny_g * nz_g + (cj + 0) * nz_g + (ck + 0)] + b1_node = flex_nodebodyid[nstart + (ci + 0) * ny_g * nz_g + (cj + 0) * nz_g + (ck + 1)] + b2_node = flex_nodebodyid[nstart + (ci + 0) * ny_g * nz_g + (cj + 1) * nz_g + (ck + 0)] + b3_node = flex_nodebodyid[nstart + (ci + 0) * ny_g * nz_g + (cj + 1) * nz_g + (ck + 1)] + b4_node = flex_nodebodyid[nstart + (ci + 1) * ny_g * nz_g + (cj + 0) * nz_g + (ck + 0)] + b5_node = flex_nodebodyid[nstart + (ci + 1) * ny_g * nz_g + (cj + 0) * nz_g + (ck + 1)] + b6_node = flex_nodebodyid[nstart + (ci + 1) * ny_g * nz_g + (cj + 1) * nz_g + (ck + 0)] + b7_node = flex_nodebodyid[nstart + (ci + 1) * ny_g * nz_g + (cj + 1) * nz_g + (ck + 1)] + + W_mat = wp.mat33(W0, W1, W2, W3, W4, W5, W6, W7, -1.0) + b_mat = wp.mat33( + float(b0_node), + float(b1_node), + float(b2_node), + float(b3_node), + float(b4_node), + float(b5_node), + float(b6_node), + float(b7_node), + -1.0, + ) + return support.select_top4_weights(W_mat, b_mat) else: return wp.vec4i(-1, -1, -1, -1), wp.vec4(0.0, 0.0, 0.0, 0.0) +# TODO(flex): Add interpolated flex self-collision support. Currently, interpolated +# flex contacts only handle flex-vs-geom contacts (e.g., flex on ground plane). +# Self-collision requires distributing both sides of the contact across cell nodes. + +# TODO(flex): Replace inverse-distance contact weights with barycentric weights +# (mj_elemBodyWeight-style) for element contacts on interpolated flex. + + @cache_kernel def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool, newton: bool): IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC @@ -2242,6 +2672,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool, newton: bool): def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: bool): IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC IS_SPARSE = is_sparse + HAS_FLEX = True @wp.kernel(module="unique", enable_backward=False) def kernel( @@ -2253,12 +2684,17 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b dof_parentid: wp.array[int], geom_bodyid: wp.array[int], flex_dim: wp.array[int], + flex_interp: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], flex_vertadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_shelldataadr: wp.array[int], + flex_nodebodyid: wp.array[int], flex_vertbodyid: wp.array[int], flex_elem: wp.array[int], flex_shell: wp.array[int], + flex_vert0: wp.array[wp.vec3], # Data in: flexvert_xpos_in: wp.array2d[wp.vec3], njmax_in: int, @@ -2266,7 +2702,6 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b nacon_in: wp.array[int], # In: dist_in: wp.array[float], - pos_in: wp.array[wp.vec3], condim_in: wp.array[int], includemargin_in: wp.array[float], worldid_in: wp.array[int], @@ -2274,6 +2709,7 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b flex_in: wp.array[wp.vec2i], elem_in: wp.array[wp.vec2i], vert_in: wp.array[wp.vec2i], + pos_in: wp.array[wp.vec3], type_in: wp.array[int], # Data out: nefc_out: wp.array[int], @@ -2341,12 +2777,16 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b body_ids1, weights1 = _get_contact_bodies_and_weights( geom_bodyid, flex_dim, + flex_cellnum, + flex_nodeadr, flex_vertadr, flex_elemdataadr, flex_shelldataadr, + flex_nodebodyid, flex_vertbodyid, flex_elem, flex_shell, + flex_vert0, flexvert_xpos_in, conid, 0, @@ -2360,12 +2800,16 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b body_ids2, weights2 = _get_contact_bodies_and_weights( geom_bodyid, flex_dim, + flex_cellnum, + flex_nodeadr, flex_vertadr, flex_elemdataadr, flex_shelldataadr, + flex_nodebodyid, flex_vertbodyid, flex_elem, flex_shell, + flex_vert0, flexvert_xpos_in, conid, 1, @@ -2377,54 +2821,168 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b worldid, ) - b1_0 = body_weldid[body_ids1[0]] - b1_1 = body_weldid[body_ids1[1]] if body_ids1[1] >= 0 else -1 - b1_2 = body_weldid[body_ids1[2]] if body_ids1[2] >= 0 else -1 - b1_3 = body_weldid[body_ids1[3]] if body_ids1[3] >= 0 else -1 + is_interp = False + if wp.static(HAS_FLEX): + if geom[0] < 0 and flex[0] >= 0 and (vert[0] >= 0 or elem[0] >= 0): + if flex_interp[flex[0]] != 0: + is_interp = True + if geom[1] < 0 and flex[1] >= 0 and (vert[1] >= 0 or elem[1] >= 0): + if flex_interp[flex[1]] != 0: + is_interp = True - b2_0 = body_weldid[body_ids2[0]] - b2_1 = body_weldid[body_ids2[1]] if body_ids2[1] >= 0 else -1 - b2_2 = body_weldid[body_ids2[2]] if body_ids2[2] >= 0 else -1 - b2_3 = body_weldid[body_ids2[3]] if body_ids2[3] >= 0 else -1 + if is_interp: + # Interpolated flex: sum of all contributing body dofnums + rownnz = int(0) + for side in range(2): + if geom[side] >= 0: + b = body_weldid[geom_bodyid[geom[side]]] + rownnz += body_dofnum[b] + elif flex[side] >= 0 and vert[side] >= 0: + f = flex[side] + if flex_interp[f] != 0: + # Compute parametric coordinate from flex_vert0 + v_adr = flex_vertadr[f] + vert[side] + coord = flex_vert0[v_adr] + cn = flex_cellnum[f] + cx = cn[0] + cy = cn[1] + cz = cn[2] - dof1_0 = int(body_dofadr[b1_0] + body_dofnum[b1_0] - 1) if b1_0 >= 0 else -1 - dof1_1 = int(body_dofadr[b1_1] + body_dofnum[b1_1] - 1) if b1_1 >= 0 else -1 - dof1_2 = int(body_dofadr[b1_2] + body_dofnum[b1_2] - 1) if b1_2 >= 0 else -1 - dof1_3 = int(body_dofadr[b1_3] + body_dofnum[b1_3] - 1) if b1_3 >= 0 else -1 + # Cell lookup + ci = wp.min(int(coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) - dof2_0 = int(body_dofadr[b2_0] + body_dofnum[b2_0] - 1) if b2_0 >= 0 else -1 - dof2_1 = int(body_dofadr[b2_1] + body_dofnum[b2_1] - 1) if b2_1 >= 0 else -1 - dof2_2 = int(body_dofadr[b2_2] + body_dofnum[b2_2] - 1) if b2_2 >= 0 else -1 - dof2_3 = int(body_dofadr[b2_3] + body_dofnum[b2_3] - 1) if b2_3 >= 0 else -1 + # Local parametric coordinates + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) - # count non-zeros - rownnz = int(0) - while ( - dof1_0 >= 0 or dof1_1 >= 0 or dof1_2 >= 0 or dof1_3 >= 0 or dof2_0 >= 0 or dof2_1 >= 0 or dof2_2 >= 0 or dof2_3 >= 0 - ): - da1_max = wp.max(dof1_0, wp.max(dof1_1, wp.max(dof1_2, dof1_3))) - da2_max = wp.max(dof2_0, wp.max(dof2_1, wp.max(dof2_2, dof2_3))) - da = wp.max(da1_max, da2_max) + # Node grid dimensions + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f] - if dof1_0 == da: - dof1_0 = dof_parentid[dof1_0] - if dof1_1 == da: - dof1_1 = dof_parentid[dof1_1] - if dof1_2 == da: - dof1_2 = dof_parentid[dof1_2] - if dof1_3 == da: - dof1_3 = dof_parentid[dof1_3] + # Loop over 8 trilinear nodes + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = body_weldid[flex_nodebodyid[nstart + node_idx]] + rownnz += body_dofnum[b] + else: + b = body_weldid[flex_vertbodyid[flex_vertadr[f] + vert[side]]] + rownnz += body_dofnum[b] + elif flex[side] >= 0 and elem[side] >= 0: + # Elem contact: use blended coordinate from distance weighting + f = flex[side] + e = elem[side] + dim_f = flex_dim[f] + edata_adr = flex_elemdataadr[f] + e * (dim_f + 1) + vert_adr_f = flex_vertadr[f] + contact_pos = pos_in[conid] - if dof2_0 == da: - dof2_0 = dof_parentid[dof2_0] - if dof2_1 == da: - dof2_1 = dof_parentid[dof2_1] - if dof2_2 == da: - dof2_2 = dof_parentid[dof2_2] - if dof2_3 == da: - dof2_3 = dof_parentid[dof2_3] + total_inv_dist = float(0.0) + blended_coord = wp.vec3(0.0, 0.0, 0.0) + for vi in range(4): + if vi <= dim_f: + v_idx = flex_elem[edata_adr + vi] + vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx] + dist_v = wp.length(contact_pos - vpos) + w_inv = 1.0 / wp.max(1.0e-10, dist_v) + total_inv_dist += w_inv + blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv + if total_inv_dist > 1.0e-10: + blended_coord = blended_coord / total_inv_dist - rownnz += 1 + if flex_interp[f] != 0: + cn = flex_cellnum[f] + cx = cn[0] + cy = cn[1] + cz = cn[2] + ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f] + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = body_weldid[flex_nodebodyid[nstart + node_idx]] + rownnz += body_dofnum[b] + else: + # Standard path (including elements up to 4 bodies) + b1_0 = body_weldid[body_ids1[0]] + b1_1 = body_weldid[body_ids1[1]] if body_ids1[1] >= 0 else -1 + b1_2 = body_weldid[body_ids1[2]] if body_ids1[2] >= 0 else -1 + b1_3 = body_weldid[body_ids1[3]] if body_ids1[3] >= 0 else -1 + + b2_0 = body_weldid[body_ids2[0]] + b2_1 = body_weldid[body_ids2[1]] if body_ids2[1] >= 0 else -1 + b2_2 = body_weldid[body_ids2[2]] if body_ids2[2] >= 0 else -1 + b2_3 = body_weldid[body_ids2[3]] if body_ids2[3] >= 0 else -1 + + dof1_0 = int(body_dofadr[b1_0] + body_dofnum[b1_0] - 1) if b1_0 >= 0 else -1 + dof1_1 = int(body_dofadr[b1_1] + body_dofnum[b1_1] - 1) if b1_1 >= 0 else -1 + dof1_2 = int(body_dofadr[b1_2] + body_dofnum[b1_2] - 1) if b1_2 >= 0 else -1 + dof1_3 = int(body_dofadr[b1_3] + body_dofnum[b1_3] - 1) if b1_3 >= 0 else -1 + + dof2_0 = int(body_dofadr[b2_0] + body_dofnum[b2_0] - 1) if b2_0 >= 0 else -1 + dof2_1 = int(body_dofadr[b2_1] + body_dofnum[b2_1] - 1) if b2_1 >= 0 else -1 + dof2_2 = int(body_dofadr[b2_2] + body_dofnum[b2_2] - 1) if b2_2 >= 0 else -1 + dof2_3 = int(body_dofadr[b2_3] + body_dofnum[b2_3] - 1) if b2_3 >= 0 else -1 + + # count non-zeros + rownnz = int(0) + while ( + dof1_0 >= 0 or dof1_1 >= 0 or dof1_2 >= 0 or dof1_3 >= 0 or dof2_0 >= 0 or dof2_1 >= 0 or dof2_2 >= 0 or dof2_3 >= 0 + ): + da1_max = wp.max(dof1_0, wp.max(dof1_1, wp.max(dof1_2, dof1_3))) + da2_max = wp.max(dof2_0, wp.max(dof2_1, wp.max(dof2_2, dof2_3))) + da = wp.max(da1_max, da2_max) + + if dof1_0 == da: + dof1_0 = dof_parentid[dof1_0] + if dof1_1 == da: + dof1_1 = dof_parentid[dof1_1] + if dof1_2 == da: + dof1_2 = dof_parentid[dof1_2] + if dof1_3 == da: + dof1_3 = dof_parentid[dof1_3] + + if dof2_0 == da: + dof2_0 = dof_parentid[dof2_0] + if dof2_1 == da: + dof2_1 = dof_parentid[dof2_1] + if dof2_2 == da: + dof2_2 = dof_parentid[dof2_2] + if dof2_3 == da: + dof2_3 = dof_parentid[dof2_3] + + rownnz += 1 rowadr = wp.atomic_add(efc_nnz_out, worldid, rownnz * ndim) if rowadr + rownnz * ndim > njmax_nnz_in: @@ -2593,6 +3151,7 @@ def _efc_contact_jac_sparse(cone_type: types.ConeType): @cache_kernel def _efc_contact_jac_sparse_flex(cone_type: types.ConeType): IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC + HAS_FLEX = True @wp.kernel(module="unique", enable_backward=False) def kernel( @@ -2606,12 +3165,17 @@ def _efc_contact_jac_sparse_flex(cone_type: types.ConeType): dof_parentid: wp.array[int], geom_bodyid: wp.array[int], flex_dim: wp.array[int], + flex_interp: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], flex_vertadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_shelldataadr: wp.array[int], + flex_nodebodyid: wp.array[int], flex_vertbodyid: wp.array[int], flex_elem: wp.array[int], flex_shell: wp.array[int], + flex_vert0: wp.array[wp.vec3], body_isdofancestor: wp.array2d[int], # Data in: qvel_in: wp.array2d[float], @@ -2655,55 +3219,14 @@ def _efc_contact_jac_sparse_flex(cone_type: types.ConeType): vert = vert_in[conid] con_pos = pos_in[conid] - body_ids1, weights1 = _get_contact_bodies_and_weights( - geom_bodyid, - flex_dim, - flex_vertadr, - flex_elemdataadr, - flex_shelldataadr, - flex_vertbodyid, - flex_elem, - flex_shell, - flexvert_xpos_in, - conid, - 0, - geom, - flex, - elem, - vert, - con_pos, - worldid, - ) - body_ids2, weights2 = _get_contact_bodies_and_weights( - geom_bodyid, - flex_dim, - flex_vertadr, - flex_elemdataadr, - flex_shelldataadr, - flex_vertbodyid, - flex_elem, - flex_shell, - flexvert_xpos_in, - conid, - 1, - geom, - flex, - elem, - vert, - con_pos, - worldid, - ) - - # skip fixed bodies - b1_0 = body_weldid[body_ids1[0]] - b1_1 = body_weldid[body_ids1[1]] if body_ids1[1] >= 0 else -1 - b1_2 = body_weldid[body_ids1[2]] if body_ids1[2] >= 0 else -1 - b1_3 = body_weldid[body_ids1[3]] if body_ids1[3] >= 0 else -1 - - b2_0 = body_weldid[body_ids2[0]] - b2_1 = body_weldid[body_ids2[1]] if body_ids2[1] >= 0 else -1 - b2_2 = body_weldid[body_ids2[2]] if body_ids2[2] >= 0 else -1 - b2_3 = body_weldid[body_ids2[3]] if body_ids2[3] >= 0 else -1 + is_interp = False + if wp.static(HAS_FLEX): + if geom[0] < 0 and flex[0] >= 0 and (vert[0] >= 0 or elem[0] >= 0): + if flex_interp[flex[0]] != 0: + is_interp = True + if geom[1] < 0 and flex[1] >= 0 and (vert[1] >= 0 or elem[1] >= 0): + if flex_interp[flex[1]] != 0: + is_interp = True if not wp.static(IS_ELLIPTIC): frame_0 = frame_in[conid, 0] @@ -2711,149 +3234,515 @@ def _efc_contact_jac_sparse_flex(cone_type: types.ConeType): dimid2 = dimid / 2 + 1 frii = friction_in[conid, dimid2 - 1] - dof1_0 = int(body_dofadr[b1_0] + body_dofnum[b1_0] - 1) if b1_0 >= 0 else -1 - dof1_1 = int(body_dofadr[b1_1] + body_dofnum[b1_1] - 1) if b1_1 >= 0 else -1 - dof1_2 = int(body_dofadr[b1_2] + body_dofnum[b1_2] - 1) if b1_2 >= 0 else -1 - dof1_3 = int(body_dofadr[b1_3] + body_dofnum[b1_3] - 1) if b1_3 >= 0 else -1 - da1 = wp.max(dof1_0, wp.max(dof1_1, wp.max(dof1_2, dof1_3))) - - dof2_0 = int(body_dofadr[b2_0] + body_dofnum[b2_0] - 1) if b2_0 >= 0 else -1 - dof2_1 = int(body_dofadr[b2_1] + body_dofnum[b2_1] - 1) if b2_1 >= 0 else -1 - dof2_2 = int(body_dofadr[b2_2] + body_dofnum[b2_2] - 1) if b2_2 >= 0 else -1 - dof2_3 = int(body_dofadr[b2_3] + body_dofnum[b2_3] - 1) if b2_3 >= 0 else -1 - da2 = wp.max(dof2_0, wp.max(dof2_1, wp.max(dof2_2, dof2_3))) - - da = wp.max(da1, da2) - rowadr = efc_J_rowadr_in[worldid, efcid] rownnz = efc_J_rownnz_in[worldid, efcid] - Jqvel = float(0.0) - nnz = int(0) - dofid = int(da) + if is_interp: + # Interpolated flex path: iterate over bodies per side, accumulate weighted Jacobians + nnz = int(0) + Jqvel = float(0.0) - while True: - if nnz >= rownnz: - break + for side in range(2): + sign = float(-1.0) if side == 0 else float(1.0) - if dofid == da: - jac1p = wp.vec3(0.0) - jac1r = wp.vec3(0.0) - if dof1_0 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_0, dofid, worldid - ) - jac1p += jp * weights1[0] - jac1r += jr * weights1[0] - if dof1_1 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_1, dofid, worldid - ) - jac1p += jp * weights1[1] - jac1r += jr * weights1[1] - if dof1_2 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_2, dofid, worldid - ) - jac1p += jp * weights1[2] - jac1r += jr * weights1[2] - if dof1_3 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_3, dofid, worldid - ) - jac1p += jp * weights1[3] - jac1r += jr * weights1[3] + if geom[side] >= 0: + # Geom side: single body + b = body_weldid[geom_bodyid[geom[side]]] + dof_start = body_dofadr[b] + ndof = body_dofnum[b] + for di in range(ndof): + dofid = dof_start + di + jacp, jacr = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + body_isdofancestor, + subtree_com_in, + cdof_in, + con_pos, + b, + dofid, + worldid, + ) - jac2p = wp.vec3(0.0) - jac2r = wp.vec3(0.0) - if dof2_0 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_0, dofid, worldid - ) - jac2p += jp * weights2[0] - jac2r += jr * weights2[0] - if dof2_1 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_1, dofid, worldid - ) - jac2p += jp * weights2[1] - jac2r += jr * weights2[1] - if dof2_2 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_2, dofid, worldid - ) - jac2p += jp * weights2[2] - jac2r += jr * weights2[2] - if dof2_3 == da: - jp, jr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_3, dofid, worldid - ) - jac2p += jp * weights2[3] - jac2r += jr * weights2[3] + if wp.static(IS_ELLIPTIC): + J = float(0.0) + if dimid < 3: + frame_row = frame_in[conid, dimid] + for xyz in range(3): + J += frame_row[xyz] * jacp[xyz] * sign + else: + frame_row = frame_in[conid, dimid - 3] + for xyz in range(3): + J += frame_row[xyz] * jacr[xyz] * sign + else: + J = float(0.0) + Ji = float(0.0) + for xyz in range(3): + J += frame_0[xyz] * jacp[xyz] * sign + if condim > 1: + if dimid2 < 3: + Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * sign + else: + Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * sign + if condim > 1: + if dimid % 2 == 0: + J += Ji * frii + else: + J -= Ji * frii - jacp_dif = jac2p - jac1p - jacr_dif = jac2r - jac1r + if nnz < rownnz: + sparseid = rowadr + nnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + Jqvel += J * qvel_in[worldid, dofid] + nnz += 1 - if wp.static(IS_ELLIPTIC): - J = float(0.0) - if dimid < 3: - frame_row = frame_in[conid, dimid] - for xyz in range(3): - J += frame_row[xyz] * jacp_dif[xyz] + elif flex[side] >= 0 and vert[side] >= 0: + f = flex[side] + if flex_interp[f] != 0: + # Interpolated flex side: compute cell node bodies and weights + v_adr = flex_vertadr[f] + vert[side] + coord = flex_vert0[v_adr] + cn = flex_cellnum[f] + cx = cn[0] + cy = cn[1] + cz = cn[2] + + ci = wp.min(int(coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f] + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = body_weldid[flex_nodebodyid[nstart + node_idx]] + w_sign = w * sign + + dof_start = body_dofadr[b] + ndof = body_dofnum[b] + for di in range(ndof): + dofid = dof_start + di + jacp, jacr = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + body_isdofancestor, + subtree_com_in, + cdof_in, + con_pos, + b, + dofid, + worldid, + ) + + if wp.static(IS_ELLIPTIC): + J = float(0.0) + if dimid < 3: + frame_row = frame_in[conid, dimid] + for xyz in range(3): + J += frame_row[xyz] * jacp[xyz] * w_sign + else: + frame_row = frame_in[conid, dimid - 3] + for xyz in range(3): + J += frame_row[xyz] * jacr[xyz] * w_sign + else: + J = float(0.0) + Ji = float(0.0) + for xyz in range(3): + J += frame_0[xyz] * jacp[xyz] * w_sign + if condim > 1: + if dimid2 < 3: + Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * w_sign + else: + Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * w_sign + if condim > 1: + if dimid % 2 == 0: + J += Ji * frii + else: + J -= Ji * frii + + if nnz < rownnz: + sparseid = rowadr + nnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + Jqvel += J * qvel_in[worldid, dofid] + nnz += 1 else: - frame_row = frame_in[conid, dimid - 3] - for xyz in range(3): - J += frame_row[xyz] * jacr_dif[xyz] - else: - J = float(0.0) - Ji = float(0.0) + # Non-interpolated flex: single body + b = body_weldid[flex_vertbodyid[flex_vertadr[f] + vert[side]]] + dof_start = body_dofadr[b] + ndof = body_dofnum[b] + for di in range(ndof): + dofid = dof_start + di + jacp, jacr = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + body_isdofancestor, + subtree_com_in, + cdof_in, + con_pos, + b, + dofid, + worldid, + ) - for xyz in range(3): - J += frame_0[xyz] * jacp_dif[xyz] + if wp.static(IS_ELLIPTIC): + J = float(0.0) + if dimid < 3: + frame_row = frame_in[conid, dimid] + for xyz in range(3): + J += frame_row[xyz] * jacp[xyz] * sign + else: + frame_row = frame_in[conid, dimid - 3] + for xyz in range(3): + J += frame_row[xyz] * jacr[xyz] * sign + else: + J = float(0.0) + Ji = float(0.0) + for xyz in range(3): + J += frame_0[xyz] * jacp[xyz] * sign + if condim > 1: + if dimid2 < 3: + Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * sign + else: + Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * sign + if condim > 1: + if dimid % 2 == 0: + J += Ji * frii + else: + J -= Ji * frii + + if nnz < rownnz: + sparseid = rowadr + nnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + Jqvel += J * qvel_in[worldid, dofid] + nnz += 1 + + elif flex[side] >= 0 and elem[side] >= 0: + # Elem contact: compute blended coordinate from distance weighting + f = flex[side] + e = elem[side] + dim_f = flex_dim[f] + edata_adr = flex_elemdataadr[f] + e * (dim_f + 1) + vert_adr_f = flex_vertadr[f] + + total_inv_dist = float(0.0) + blended_coord = wp.vec3(0.0, 0.0, 0.0) + for vi in range(4): + if vi <= dim_f: + v_idx = flex_elem[edata_adr + vi] + vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx] + dist_v = wp.length(con_pos - vpos) + w_inv = 1.0 / wp.max(1.0e-10, dist_v) + total_inv_dist += w_inv + blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv + if total_inv_dist > 1.0e-10: + blended_coord = blended_coord / total_inv_dist + + if flex_interp[f] != 0: + cn = flex_cellnum[f] + cx = cn[0] + cy = cn[1] + cz = cn[2] + ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f] + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = body_weldid[flex_nodebodyid[nstart + node_idx]] + w_sign = w * sign + + dof_start = body_dofadr[b] + ndof = body_dofnum[b] + for di in range(ndof): + dofid = dof_start + di + jacp, jacr = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + body_isdofancestor, + subtree_com_in, + cdof_in, + con_pos, + b, + dofid, + worldid, + ) + + if wp.static(IS_ELLIPTIC): + J = float(0.0) + if dimid < 3: + frame_row = frame_in[conid, dimid] + for xyz in range(3): + J += frame_row[xyz] * jacp[xyz] * w_sign + else: + frame_row = frame_in[conid, dimid - 3] + for xyz in range(3): + J += frame_row[xyz] * jacr[xyz] * w_sign + else: + J = float(0.0) + Ji = float(0.0) + for xyz in range(3): + J += frame_0[xyz] * jacp[xyz] * w_sign + if condim > 1: + if dimid2 < 3: + Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * w_sign + else: + Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * w_sign + if condim > 1: + if dimid % 2 == 0: + J += Ji * frii + else: + J -= Ji * frii + + if nnz < rownnz: + sparseid = rowadr + nnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + Jqvel += J * qvel_in[worldid, dofid] + nnz += 1 + + efc_Jqvel_out[worldid, efcid] = Jqvel + + else: + # Non-interpolated flex path: use HEAD's 4-body lockstep tree-walk + body_ids1, weights1 = _get_contact_bodies_and_weights( + geom_bodyid, + flex_dim, + flex_cellnum, + flex_nodeadr, + flex_vertadr, + flex_elemdataadr, + flex_shelldataadr, + flex_nodebodyid, + flex_vertbodyid, + flex_elem, + flex_shell, + flex_vert0, + flexvert_xpos_in, + conid, + 0, + geom, + flex, + elem, + vert, + con_pos, + worldid, + ) + body_ids2, weights2 = _get_contact_bodies_and_weights( + geom_bodyid, + flex_dim, + flex_cellnum, + flex_nodeadr, + flex_vertadr, + flex_elemdataadr, + flex_shelldataadr, + flex_nodebodyid, + flex_vertbodyid, + flex_elem, + flex_shell, + flex_vert0, + flexvert_xpos_in, + conid, + 1, + geom, + flex, + elem, + vert, + con_pos, + worldid, + ) + + # skip fixed bodies + b1_0 = body_weldid[body_ids1[0]] + b1_1 = body_weldid[body_ids1[1]] if body_ids1[1] >= 0 else -1 + b1_2 = body_weldid[body_ids1[2]] if body_ids1[2] >= 0 else -1 + b1_3 = body_weldid[body_ids1[3]] if body_ids1[3] >= 0 else -1 + + b2_0 = body_weldid[body_ids2[0]] + b2_1 = body_weldid[body_ids2[1]] if body_ids2[1] >= 0 else -1 + b2_2 = body_weldid[body_ids2[2]] if body_ids2[2] >= 0 else -1 + b2_3 = body_weldid[body_ids2[3]] if body_ids2[3] >= 0 else -1 + + dof1_0 = int(body_dofadr[b1_0] + body_dofnum[b1_0] - 1) if b1_0 >= 0 else -1 + dof1_1 = int(body_dofadr[b1_1] + body_dofnum[b1_1] - 1) if b1_1 >= 0 else -1 + dof1_2 = int(body_dofadr[b1_2] + body_dofnum[b1_2] - 1) if b1_2 >= 0 else -1 + dof1_3 = int(body_dofadr[b1_3] + body_dofnum[b1_3] - 1) if b1_3 >= 0 else -1 + da1 = wp.max(dof1_0, wp.max(dof1_1, wp.max(dof1_2, dof1_3))) + + dof2_0 = int(body_dofadr[b2_0] + body_dofnum[b2_0] - 1) if b2_0 >= 0 else -1 + dof2_1 = int(body_dofadr[b2_1] + body_dofnum[b2_1] - 1) if b2_1 >= 0 else -1 + dof2_2 = int(body_dofadr[b2_2] + body_dofnum[b2_2] - 1) if b2_2 >= 0 else -1 + dof2_3 = int(body_dofadr[b2_3] + body_dofnum[b2_3] - 1) if b2_3 >= 0 else -1 + da2 = wp.max(dof2_0, wp.max(dof2_1, wp.max(dof2_2, dof2_3))) + + da = wp.max(da1, da2) + + Jqvel = float(0.0) + nnz = int(0) + dofid = int(da) + + while True: + if nnz >= rownnz: + break + + if dofid == da: + jac1p = wp.vec3(0.0) + jac1r = wp.vec3(0.0) + if dof1_0 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_0, dofid, worldid + ) + jac1p += jp * weights1[0] + jac1r += jr * weights1[0] + if dof1_1 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_1, dofid, worldid + ) + jac1p += jp * weights1[1] + jac1r += jr * weights1[1] + if dof1_2 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_2, dofid, worldid + ) + jac1p += jp * weights1[2] + jac1r += jr * weights1[2] + if dof1_3 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b1_3, dofid, worldid + ) + jac1p += jp * weights1[3] + jac1r += jr * weights1[3] + + jac2p = wp.vec3(0.0) + jac2r = wp.vec3(0.0) + if dof2_0 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_0, dofid, worldid + ) + jac2p += jp * weights2[0] + jac2r += jr * weights2[0] + if dof2_1 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_1, dofid, worldid + ) + jac2p += jp * weights2[1] + jac2r += jr * weights2[1] + if dof2_2 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_2, dofid, worldid + ) + jac2p += jp * weights2[2] + jac2r += jr * weights2[2] + if dof2_3 == da: + jp, jr = support.jac_dof( + body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b2_3, dofid, worldid + ) + jac2p += jp * weights2[3] + jac2r += jr * weights2[3] + + jacp_dif = jac2p - jac1p + jacr_dif = jac2r - jac1r + + if wp.static(IS_ELLIPTIC): + J = float(0.0) + if dimid < 3: + frame_row = frame_in[conid, dimid] + for xyz in range(3): + J += frame_row[xyz] * jacp_dif[xyz] + else: + frame_row = frame_in[conid, dimid - 3] + for xyz in range(3): + J += frame_row[xyz] * jacr_dif[xyz] + else: + J = float(0.0) + Ji = float(0.0) + + for xyz in range(3): + J += frame_0[xyz] * jacp_dif[xyz] + + if condim > 1: + if dimid2 < 3: + Ji += frame_in[conid, dimid2][xyz] * jacp_dif[xyz] + else: + Ji += frame_in[conid, dimid2 - 3][xyz] * jacr_dif[xyz] if condim > 1: - if dimid2 < 3: - Ji += frame_in[conid, dimid2][xyz] * jacp_dif[xyz] + if dimid % 2 == 0: + J += Ji * frii else: - Ji += frame_in[conid, dimid2 - 3][xyz] * jacr_dif[xyz] + J -= Ji * frii - if condim > 1: - if dimid % 2 == 0: - J += Ji * frii - else: - J -= Ji * frii + sparseid = rowadr + nnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + nnz += 1 + Jqvel += J * qvel_in[worldid, dofid] - sparseid = rowadr + nnz - efc_J_colind_out[worldid, 0, sparseid] = dofid - efc_J_out[worldid, 0, sparseid] = J - nnz += 1 - Jqvel += J * qvel_in[worldid, dofid] + # Advance tree pointers + if dof1_0 == da: + dof1_0 = dof_parentid[dof1_0] + if dof1_1 == da: + dof1_1 = dof_parentid[dof1_1] + if dof1_2 == da: + dof1_2 = dof_parentid[dof1_2] + if dof1_3 == da: + dof1_3 = dof_parentid[dof1_3] + da1 = wp.max(dof1_0, wp.max(dof1_1, wp.max(dof1_2, dof1_3))) - # Advance tree pointers - if dof1_0 == da: - dof1_0 = dof_parentid[dof1_0] - if dof1_1 == da: - dof1_1 = dof_parentid[dof1_1] - if dof1_2 == da: - dof1_2 = dof_parentid[dof1_2] - if dof1_3 == da: - dof1_3 = dof_parentid[dof1_3] - da1 = wp.max(dof1_0, wp.max(dof1_1, wp.max(dof1_2, dof1_3))) + if dof2_0 == da: + dof2_0 = dof_parentid[dof2_0] + if dof2_1 == da: + dof2_1 = dof_parentid[dof2_1] + if dof2_2 == da: + dof2_2 = dof_parentid[dof2_2] + if dof2_3 == da: + dof2_3 = dof_parentid[dof2_3] + da2 = wp.max(dof2_0, wp.max(dof2_1, wp.max(dof2_2, dof2_3))) - if dof2_0 == da: - dof2_0 = dof_parentid[dof2_0] - if dof2_1 == da: - dof2_1 = dof_parentid[dof2_1] - if dof2_2 == da: - dof2_2 = dof_parentid[dof2_2] - if dof2_3 == da: - dof2_3 = dof_parentid[dof2_3] - da2 = wp.max(dof2_0, wp.max(dof2_1, wp.max(dof2_2, dof2_3))) + da = wp.max(da1, da2) + dofid = da - da = wp.max(da1, da2) - dofid = da - - efc_Jqvel_out[worldid, efcid] = Jqvel + efc_Jqvel_out[worldid, efcid] = Jqvel return kernel @@ -2920,17 +3809,16 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType): con_pos = pos_in[conid] offset1 = con_pos - subtree_com_in[worldid, body_rootid[body1]] - offset2 = con_pos - subtree_com_in[worldid, body_rootid[body2]] - affects1_tile = wp.tile_load(body_isdofancestor[body1], shape=TILE_SIZE, offset=dof_start, bounds_check=True) - affects2_tile = wp.tile_load(body_isdofancestor[body2], shape=TILE_SIZE, offset=dof_start, bounds_check=True) - jacp1_tile = wp.tile_map(support._compute_jacp, cdof_tile, offset1, affects1_tile) - jacp2_tile = wp.tile_map(support._compute_jacp, cdof_tile, offset2, affects2_tile) - jacp_dif_tile = wp.tile_map(wp.sub, jacp2_tile, jacp1_tile) - jacr1_tile = wp.tile_map(support._compute_jacr, cdof_tile, affects1_tile) + + offset2 = con_pos - subtree_com_in[worldid, body_rootid[body2]] + affects2_tile = wp.tile_load(body_isdofancestor[body2], shape=TILE_SIZE, offset=dof_start, bounds_check=True) + jacp2_tile = wp.tile_map(support._compute_jacp, cdof_tile, offset2, affects2_tile) jacr2_tile = wp.tile_map(support._compute_jacr, cdof_tile, affects2_tile) + + jacp_dif_tile = wp.tile_map(wp.sub, jacp2_tile, jacp1_tile) jacr_dif_tile = wp.tile_map(wp.sub, jacr2_tile, jacr1_tile) if not wp.static(IS_ELLIPTIC): @@ -2999,12 +3887,16 @@ def _efc_contact_jac_dense_flex(tile_size: int, cone_type: types.ConeType): body_rootid: wp.array[int], geom_bodyid: wp.array[int], flex_dim: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], flex_vertadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_shelldataadr: wp.array[int], + flex_nodebodyid: wp.array[int], flex_vertbodyid: wp.array[int], flex_elem: wp.array[int], flex_shell: wp.array[int], + flex_vert0: wp.array[wp.vec3], body_isdofancestor: wp.array2d[int], # Data in: ne_in: wp.array[int], @@ -3064,12 +3956,16 @@ def _efc_contact_jac_dense_flex(tile_size: int, cone_type: types.ConeType): body_ids1, weights1 = _get_contact_bodies_and_weights( geom_bodyid, flex_dim, + flex_cellnum, + flex_nodeadr, flex_vertadr, flex_elemdataadr, flex_shelldataadr, + flex_nodebodyid, flex_vertbodyid, flex_elem, flex_shell, + flex_vert0, flexvert_xpos_in, conid, 0, @@ -3083,12 +3979,16 @@ def _efc_contact_jac_dense_flex(tile_size: int, cone_type: types.ConeType): body_ids2, weights2 = _get_contact_bodies_and_weights( geom_bodyid, flex_dim, + flex_cellnum, + flex_nodeadr, flex_vertadr, flex_elemdataadr, flex_shelldataadr, + flex_nodebodyid, flex_vertbodyid, flex_elem, flex_shell, + flex_vert0, flexvert_xpos_in, conid, 1, @@ -3445,12 +4345,17 @@ def _efc_contact_update_flex(cone_type: types.ConeType): body_invweight0: wp.array2d[wp.vec2], geom_bodyid: wp.array[int], flex_dim: wp.array[int], + flex_interp: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], flex_vertadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_shelldataadr: wp.array[int], + flex_nodebodyid: wp.array[int], flex_vertbodyid: wp.array[int], flex_elem: wp.array[int], flex_shell: wp.array[int], + flex_vert0: wp.array[wp.vec3], # Data in: flexvert_xpos_in: wp.array2d[wp.vec3], contact_efc_address_in: wp.array2d[int], @@ -3518,72 +4423,268 @@ def _efc_contact_update_flex(cone_type: types.ConeType): con_pos = pos_in[conid] Jqvel = efc_Jqvel_in[worldid, efcid] - body_ids1, weights1 = _get_contact_bodies_and_weights( - geom_bodyid, - flex_dim, - flex_vertadr, - flex_elemdataadr, - flex_shelldataadr, - flex_vertbodyid, - flex_elem, - flex_shell, - flexvert_xpos_in, - conid, - 0, - geom, - flex, - elem, - vert, - con_pos, - worldid, - ) - body_ids2, weights2 = _get_contact_bodies_and_weights( - geom_bodyid, - flex_dim, - flex_vertadr, - flex_elemdataadr, - flex_shelldataadr, - flex_vertbodyid, - flex_elem, - flex_shell, - flexvert_xpos_in, - conid, - 1, - geom, - flex, - elem, - vert, - con_pos, - worldid, - ) - - b1_0 = body_ids1[0] - b1_1 = body_ids1[1] - b1_2 = body_ids1[2] - b1_3 = body_ids1[3] - - b2_0 = body_ids2[0] - b2_1 = body_ids2[1] - b2_2 = body_ids2[2] - b2_3 = body_ids2[3] - body_invweight0_id = worldid % body_invweight0.shape[0] + invweight1 = float(0.0) - invweight1 = weights1[0] * body_invweight0[body_invweight0_id, b1_0][0] - if b1_1 >= 0: - invweight1 += weights1[1] * body_invweight0[body_invweight0_id, b1_1][0] - if b1_2 >= 0: - invweight1 += weights1[2] * body_invweight0[body_invweight0_id, b1_2][0] - if b1_3 >= 0: - invweight1 += weights1[3] * body_invweight0[body_invweight0_id, b1_3][0] + if geom[0] >= 0: + b1 = geom_bodyid[geom[0]] + invweight1 = body_invweight0[body_invweight0_id, b1][0] + elif flex[0] >= 0: + f1 = flex[0] + if flex_interp[f1] != 0: + if vert[0] >= 0: + v_adr = flex_vertadr[f1] + vert[0] + coord = flex_vert0[v_adr] + cn = flex_cellnum[f1] + cx = cn[0] + cy = cn[1] + cz = cn[2] - invweight2 = weights2[0] * body_invweight0[body_invweight0_id, b2_0][0] - if b2_1 >= 0: - invweight2 += weights2[1] * body_invweight0[body_invweight0_id, b2_1][0] - if b2_2 >= 0: - invweight2 += weights2[2] * body_invweight0[body_invweight0_id, b2_2][0] - if b2_3 >= 0: - invweight2 += weights2[3] * body_invweight0[body_invweight0_id, b2_3][0] + ci = wp.min(int(coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f1] + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = flex_nodebodyid[nstart + node_idx] + invweight1 += body_invweight0[body_invweight0_id, b][0] * w + elif elem[0] >= 0: + e = elem[0] + dim_f = flex_dim[f1] + edata_adr = flex_elemdataadr[f1] + e * (dim_f + 1) + vert_adr_f = flex_vertadr[f1] + + total_inv_dist = float(0.0) + blended_coord = wp.vec3(0.0, 0.0, 0.0) + for vi in range(4): + if vi <= dim_f: + v_idx = flex_elem[edata_adr + vi] + vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx] + dist_v = wp.length(con_pos - vpos) + w_inv = 1.0 / wp.max(1.0e-10, dist_v) + total_inv_dist += w_inv + blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv + + if total_inv_dist > 1.0e-10: + blended_coord = blended_coord / total_inv_dist + + cn = flex_cellnum[f1] + cx = cn[0] + cy = cn[1] + cz = cn[2] + + ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f1] + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = flex_nodebodyid[nstart + node_idx] + invweight1 += body_invweight0[body_invweight0_id, b][0] * w + else: + body_ids, weights = _get_contact_bodies_and_weights( + geom_bodyid, + flex_dim, + flex_cellnum, + flex_nodeadr, + flex_vertadr, + flex_elemdataadr, + flex_shelldataadr, + flex_nodebodyid, + flex_vertbodyid, + flex_elem, + flex_shell, + flex_vert0, + flexvert_xpos_in, + conid, + 0, + geom, + flex, + elem, + vert, + con_pos, + worldid, + ) + b0 = body_ids[0] + b1 = body_ids[1] + b2 = body_ids[2] + b3 = body_ids[3] + invweight1 = weights[0] * body_invweight0[body_invweight0_id, b0][0] + if b1 >= 0: + invweight1 += weights[1] * body_invweight0[body_invweight0_id, b1][0] + if b2 >= 0: + invweight1 += weights[2] * body_invweight0[body_invweight0_id, b2][0] + if b3 >= 0: + invweight1 += weights[3] * body_invweight0[body_invweight0_id, b3][0] + + invweight2 = float(0.0) + + if geom[1] >= 0: + b2 = geom_bodyid[geom[1]] + invweight2 = body_invweight0[body_invweight0_id, b2][0] + elif flex[1] >= 0: + f2 = flex[1] + if flex_interp[f2] != 0: + if vert[1] >= 0: + v_adr = flex_vertadr[f2] + vert[1] + coord = flex_vert0[v_adr] + cn = flex_cellnum[f2] + cx = cn[0] + cy = cn[1] + cz = cn[2] + + ci = wp.min(int(coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f2] + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = flex_nodebodyid[nstart + node_idx] + invweight2 += body_invweight0[body_invweight0_id, b][0] * w + elif elem[1] >= 0: + e = elem[1] + dim_f = flex_dim[f2] + edata_adr = flex_elemdataadr[f2] + e * (dim_f + 1) + vert_adr_f = flex_vertadr[f2] + + total_inv_dist = float(0.0) + blended_coord = wp.vec3(0.0, 0.0, 0.0) + for vi in range(4): + if vi <= dim_f: + v_idx = flex_elem[edata_adr + vi] + vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx] + dist_v = wp.length(con_pos - vpos) + w_inv = 1.0 / wp.max(1.0e-10, dist_v) + total_inv_dist += w_inv + blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv + + if total_inv_dist > 1.0e-10: + blended_coord = blended_coord / total_inv_dist + + cn = flex_cellnum[f2] + cx = cn[0] + cy = cn[1] + cz = cn[2] + + ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f2] + + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + if w > 1.0e-5: + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + b = flex_nodebodyid[nstart + node_idx] + invweight2 += body_invweight0[body_invweight0_id, b][0] * w + else: + body_ids, weights = _get_contact_bodies_and_weights( + geom_bodyid, + flex_dim, + flex_cellnum, + flex_nodeadr, + flex_vertadr, + flex_elemdataadr, + flex_shelldataadr, + flex_nodebodyid, + flex_vertbodyid, + flex_elem, + flex_shell, + flex_vert0, + flexvert_xpos_in, + conid, + 1, + geom, + flex, + elem, + vert, + con_pos, + worldid, + ) + b0 = body_ids[0] + b1 = body_ids[1] + b2 = body_ids[2] + b3 = body_ids[3] + invweight2 = weights[0] * body_invweight0[body_invweight0_id, b0][0] + if b1 >= 0: + invweight2 += weights[1] * body_invweight0[body_invweight0_id, b1][0] + if b2 >= 0: + invweight2 += weights[2] * body_invweight0[body_invweight0_id, b2][0] + if b3 >= 0: + invweight2 += weights[3] * body_invweight0[body_invweight0_id, b3][0] invweight = invweight1 + invweight2 @@ -3881,52 +4982,119 @@ def make_constraint(m: types.Model, d: types.Data): ], ) - wp.launch( - _equality_flex(m.is_sparse, newton), - dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge), - inputs=[ - m.nv, - m.opt.timestep, - m.opt.disableflags, - m.flex_edgeadr, - m.flex_edgenum, - m.flexedge_length0, - m.flexedge_invweight0, - m.flexedge_J_rownnz, - m.flexedge_J_rowadr, - m.flexedge_J_colind, - m.eq_obj1id, - m.eq_solref, - m.eq_solimp, - m.eq_flex_adr, - d.qvel, - d.eq_active, - d.flexedge_J, - d.flexedge_length, - d.njmax, - d.njmax_nnz, - ], - outputs=[ - d.ne, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.jtdaj_adr, - d.efc.jtdaj_nrow, - d.efc.jtdaj_nblock, - d.efc.J_rownnz, - d.efc.J_rowadr, - d.efc.J_colind, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - efc_nnz, - ], - ) + if m.nflex > 0: + wp.launch( + _equality_flex(m.is_sparse, newton), + dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.flex_interp, + m.flex_edgeadr, + m.flex_edgenum, + m.flexedge_length0, + m.flexedge_invweight0, + m.flexedge_J_rownnz, + m.flexedge_J_rowadr, + m.flexedge_J_colind, + m.eq_obj1id, + m.eq_solref, + m.eq_solimp, + m.eq_flex_adr, + d.qvel, + d.eq_active, + d.flexedge_J, + d.flexedge_length, + d.njmax, + d.njmax_nnz, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.jtdaj_adr, + d.efc.jtdaj_nrow, + d.efc.jtdaj_nblock, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + efc_nnz, + ], + ) + + if m.eq_flexstrain_adr.size: + wp.launch( + _equality_flexstrain(m.is_sparse, newton), + dim=(d.nworld, m.eq_flexstrain_adr.size), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.body_parentid, + m.body_rootid, + m.body_dofnum, + m.body_dofadr, + m.body_invweight0, + m.dof_bodyid, + m.flex_interp, + m.flex_cellnum, + m.flex_nodeadr, + m.flex_nodenum, + m.flex_stiffnessadr, + m.flex_nodebodyid, + m.flex_node, + m.flex_node0, + m.flex_stiffness, + m.flex_centered, + m.eq_obj1id, + m.eq_solref, + m.eq_solimp, + m.eq_data, + m.body_isdofancestor, + m.eq_flexstrain_adr, + m.flexstrain_J_rownnz, + m.flexstrain_J_rowadr, + m.flexstrain_J_colind, + d.qvel, + d.eq_active, + d.xpos, + d.xmat, + d.subtree_com, + d.cdof, + d.njmax, + d.njmax_nnz, + d.flexnode_xpos, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.jtdaj_adr, + d.efc.jtdaj_nrow, + d.efc.jtdaj_nblock, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + efc_nnz, + ], + ) if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS): wp.launch( @@ -4158,7 +5326,6 @@ def make_constraint(m: types.Model, d: types.Data): ) has_flex = m.nflex > 0 - if has_flex: wp.launch( _efc_contact_init_flex(m.opt.cone, m.is_sparse, newton), @@ -4171,18 +5338,22 @@ def make_constraint(m: types.Model, d: types.Data): m.dof_parentid, m.geom_bodyid, m.flex_dim, + m.flex_interp, + m.flex_cellnum, + m.flex_nodeadr, m.flex_vertadr, m.flex_elemdataadr, m.flex_shelldataadr, + m.flex_nodebodyid, m.flex_vertbodyid, m.flex_elem, m.flex_shell, + m.flex_vert0, d.flexvert_xpos, d.njmax, d.njmax_nnz, d.nacon, d.contact.dist, - d.contact.pos, d.contact.dim, d.contact.includemargin, d.contact.worldid, @@ -4190,6 +5361,7 @@ def make_constraint(m: types.Model, d: types.Data): d.contact.flex, d.contact.elem, d.contact.vert, + d.contact.pos, d.contact.type, ], outputs=[ @@ -4252,12 +5424,17 @@ def make_constraint(m: types.Model, d: types.Data): m.dof_parentid, m.geom_bodyid, m.flex_dim, + m.flex_interp, + m.flex_cellnum, + m.flex_nodeadr, m.flex_vertadr, m.flex_elemdataadr, m.flex_shelldataadr, + m.flex_nodebodyid, m.flex_vertbodyid, m.flex_elem, m.flex_shell, + m.flex_vert0, m.body_isdofancestor, d.qvel, d.subtree_com, @@ -4330,12 +5507,16 @@ def make_constraint(m: types.Model, d: types.Data): m.body_rootid, m.geom_bodyid, m.flex_dim, + m.flex_cellnum, + m.flex_nodeadr, m.flex_vertadr, m.flex_elemdataadr, m.flex_shelldataadr, + m.flex_nodebodyid, m.flex_vertbodyid, m.flex_elem, m.flex_shell, + m.flex_vert0, m.body_isdofancestor, d.ne, d.nf, @@ -4407,12 +5588,17 @@ def make_constraint(m: types.Model, d: types.Data): m.body_invweight0, m.geom_bodyid, m.flex_dim, + m.flex_interp, + m.flex_cellnum, + m.flex_nodeadr, m.flex_vertadr, m.flex_elemdataadr, m.flex_shelldataadr, + m.flex_nodebodyid, m.flex_vertbodyid, m.flex_elem, m.flex_shell, + m.flex_vert0, d.flexvert_xpos, d.contact.efc_address, d.efc.Jqvel, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index 9a2fd58e..090a3fc2 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -143,6 +143,20 @@ def _jtdaj_groups(mjd: mujoco.MjData) -> tuple[np.ndarray, np.ndarray]: return adr, nrow +def _get_nflexintcell(mjm: mujoco.MjModel) -> int: + nflexintcell = 0 + if mjm.nflex > 0 and hasattr(mjm, "flex_interp"): + for fi in range(mjm.nflex): + order = abs(int(mjm.flex_interp[fi])) + if order == 0: + continue + if hasattr(mjm, "flex_edgeequality") and mjm.flex_edgeequality[fi] == 3: + continue + cx, cy, cz = mjm.flex_cellnum[fi] + nflexintcell += int(cx) * int(cy) * int(cz) + return nflexintcell + + def is_sparse(mjm: mujoco.MjModel) -> bool: if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO: if mjm.nv > 32: @@ -330,6 +344,9 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> if (mjm.opt.enableflags & mujoco.mjtEnableBit.mjENBL_SLEEP) and (mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX).any(): raise NotImplementedError("Flex equality constraints are not supported with sleeping enabled.") + if mjm.nflex > 0 and (mjm.flex_interp < 0).any(): + raise NotImplementedError("Flex interpolation order < 0 (shell/quad elements) is not supported.") + if mjm.opt.noslip_iterations > 0: raise NotImplementedError(f"noslip solver not implemented.") @@ -437,10 +454,16 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> condim_arrays = [np.array([0]), mjm.geom_condim, mjm.pair_dim] if mjm.nflex > 0: condim_arrays.append(mjm.flex_condim) + if (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any(): + raise NotImplementedError("Flex-SDF collision is not implemented.") + if (mjm.geom_type == mujoco.mjtGeom.mjGEOM_HFIELD).any(): + raise NotImplementedError("Flex-HField collision is not implemented.") m.nmaxcondim = np.concatenate(condim_arrays).max() m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1)) m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any() + m.has_ellipsoid_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID).any() m.has_flex_selfcollide = bool(mjm.nflex > 0 and np.any(mjm.flex_selfcollide != 0)) + m.has_3d_flex = bool(mjm.nflex > 0 and np.any(mjm.flex_dim == 3)) m.max_flex_dim = int(np.max(mjm.flex_dim)) if mjm.nflex > 0 else 0 m.block_dim = types.BlockDim() # Derive CG solver block_dim from nv: clamp(round_up_to_32(nv), 32, 256) @@ -453,6 +476,28 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> m.block_dim.linesearch_iterative = 512 m.is_sparse = is_sparse(mjm) m.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0 + + m.nflexintcell = _get_nflexintcell(mjm) + + # Precompute flex_cell_map + flex_cell_map = [] + if mjm.nflex > 0 and hasattr(mjm, "flex_interp"): + for fi in range(mjm.nflex): + order = abs(int(mjm.flex_interp[fi])) + if order == 0: + continue + if hasattr(mjm, "flex_edgeequality") and mjm.flex_edgeequality[fi] == 3: + continue + cx, cy, cz = mjm.flex_cellnum[fi] + for ci in range(cx): + for cj in range(cy): + for ck in range(cz): + flex_cell_map.append((fi, ci, cj, ck)) + if not flex_cell_map: + m.flex_cell_map = np.zeros((0, 4), dtype=np.int32) + else: + m.flex_cell_map = np.array(flex_cell_map, dtype=np.int32) + m.max_ten_J_rownnz = int(mjm.ten_J_rownnz.max()) if mjm.ntendon else 0 # body ids grouped by tree level (depth-based traversal) @@ -699,6 +744,49 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> m.eq_jnt_adr = np.nonzero(mjm.eq_type == types.EqType.JOINT)[0] m.eq_ten_adr = np.nonzero(mjm.eq_type == types.EqType.TENDON)[0] m.eq_flex_adr = np.nonzero(mjm.eq_type == types.EqType.FLEX)[0] + m.eq_flexstrain_adr = np.nonzero(mjm.eq_type == types.EqType.FLEXSTRAIN)[0] + m.neq_flexstrain = m.eq_flexstrain_adr.size + + # Precompute flex strain Jacobian sparsity pattern + flexstrain_J_rownnz = [] + flexstrain_J_colind = [] + + if m.neq_flexstrain > 0: + for eqstrainid, eqid in enumerate(m.eq_flexstrain_adr): + f = int(mjm.eq_obj1id[eqid]) + order = int(mjm.flex_interp[f]) + ci = int(mjm.eq_data[eqid, 0]) + cj = int(mjm.eq_data[eqid, 1]) + ck = int(mjm.eq_data[eqid, 2]) + + cellnum = mjm.flex_cellnum[f] + cy = cellnum[1] + cz = cellnum[2] + nstart = mjm.flex_nodeadr[f] + ny_g = cy * order + 1 + nz_g = cz * order + 1 + + node_bodies = [ + mjm.flex_nodebodyid[nstart + (ci * order + li) * ny_g * nz_g + (cj * order + lj) * nz_g + (ck * order + lk)] + for li in range(order + 1) + for lj in range(order + 1) + for lk in range(order + 1) + ] + + active_dof_mask = np.any(body_isdofancestor[node_bodies, :] != 0, axis=0) + sorted_dofs = np.nonzero(active_dof_mask)[0].tolist() + flexstrain_J_rownnz.append(len(sorted_dofs)) + flexstrain_J_colind.extend(sorted_dofs) + + m.flexstrain_J_rownnz = np.array(flexstrain_J_rownnz, dtype=np.int32) + m.flexstrain_J_colind = np.array(flexstrain_J_colind, dtype=np.int32) + m.flexstrain_J_rowadr = np.cumsum([0] + flexstrain_J_rownnz[:-1], dtype=np.int32) + else: + m.flexstrain_J_rownnz = np.zeros((0,), dtype=np.int32) + m.flexstrain_J_rowadr = np.zeros((0,), dtype=np.int32) + m.flexstrain_J_colind = np.zeros((0,), dtype=np.int32) + + m.nJfs = m.flexstrain_J_colind.size # fixed tendon m.tendon_jnt_adr, m.wrap_jnt_adr = [], [] @@ -952,7 +1040,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> # Populate lookup maps and candidate pairs flexelem_geom_pairs = [] - flexshell_geom_pairs = [] flexvert_geom_pairs = [] flex_elemflexid = np.zeros(mjm.nflexelem, dtype=np.int32) @@ -996,6 +1083,7 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> mujoco.mjtGeom.mjGEOM_BOX, mujoco.mjtGeom.mjGEOM_CYLINDER, mujoco.mjtGeom.mjGEOM_MESH, + mujoco.mjtGeom.mjGEOM_ELLIPSOID, ], ) is_pl = mjm.geom_type == mujoco.mjtGeom.mjGEOM_PLANE @@ -1018,19 +1106,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> for g in matching_primitive_geoms[~filtered]: flexelem_geom_pairs.append((elemid, g)) - elif fdim == 3: - shelldata_start = mjm.flex_shelldataadr[fi] - prev_shells_offset = shell_offset - shell_num - for s in range(shell_num): - v0 = vert_start + mjm.flex_shell[shelldata_start + s * 3] - v1 = vert_start + mjm.flex_shell[shelldata_start + s * 3 + 1] - v2 = vert_start + mjm.flex_shell[shelldata_start + s * 3 + 2] - - if len(matching_primitive_geoms) > 0: - filtered = _filter_tri_geoms(mjm, v0, v1, v2, matching_primitive_geoms, filterparent) - for g in matching_primitive_geoms[~filtered]: - flexshell_geom_pairs.append((prev_shells_offset + s, g)) - # Planes vs Vertices if len(matching_plane_geoms) > 0: vert_count = mjm.flex_vertnum[fi] @@ -1057,13 +1132,10 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) -> if not flexelem_geom_pairs: flexelem_geom_pairs = np.zeros((0, 2), dtype=np.int32) - if not flexshell_geom_pairs: - flexshell_geom_pairs = np.zeros((0, 2), dtype=np.int32) if not flexvert_geom_pairs: flexvert_geom_pairs = np.zeros((0, 2), dtype=np.int32) m.flexelem_geom_pair_filtered = np.array(flexelem_geom_pairs, dtype=np.int32) - m.flexshell_geom_pair_filtered = np.array(flexshell_geom_pairs, dtype=np.int32) m.flexvert_geom_pair_filtered = np.array(flexvert_geom_pairs, dtype=np.int32) m.flex_elemflexid = flex_elemflexid @@ -1208,6 +1280,15 @@ def _default_njmax_nnz(mjm: mujoco.MjModel, nconmax: int, njmax: int) -> int: for e in range(edge_count): total_nnz += mjm.flexedge_J_rownnz[edge_start + e] + elif eq_type == mujoco.mjtEq.mjEQ_FLEXSTRAIN: + # strain constraints: each cell produces neig rows, each dense (nv) + obj1id = mjm.eq_obj1id[i] + if obj1id < mjm.nflex and hasattr(mjm, "flex_stiffnessadr"): + # estimate neig from stiffness data + adr = mjm.flex_stiffnessadr[obj1id] + neig = int(mjm.flex_stiffness[adr]) + total_nnz += neig * mjm.nv + # friction constraints total_nnz += (mjm.dof_frictionloss > 0).sum() for i in range(mjm.ntendon): @@ -1494,6 +1575,7 @@ def make_data( sizes["njmax"] = njmax sizes["nvmax"] = nvmax sizes["nvmax_pad"] = _nvmax_pad(nvmax) + sizes["nflexintcell"] = _get_nflexintcell(mjm) if njmax_nnz is None: if is_sparse(mjm): @@ -2523,6 +2605,83 @@ def _copy_tendon_length0( tendon_length0_out[tendon_length0_id, tenid] = ten_length_in[worldid, tenid] +@wp.kernel +def _compute_eq_data0( + # Model: + eq_type: wp.array[int], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], + # Data in: + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], + xmat_in: wp.array2d[wp.mat33], + # Out: + eq_data_out: wp.array2d[types.vec11], +): + """Compute eq_data for connect/weld constraints. + + Kinematics must have been evaluated at qpos0 so the constraint is satisfied at qpos0. + """ + worldid, eqid = wp.tid() + eq_data_id = worldid % eq_data_out.shape[0] + + eqtype = eq_type[eqid] + objtype = eq_objtype[eqid] + data = eq_data_out[eq_data_id, eqid] + + if eqtype == int(types.EqType.CONNECT.value): + if objtype == int(types.ObjType.BODY.value): + obj1id = eq_obj1id[eqid] + obj2id = eq_obj2id[eqid] + + # data[0:3] = anchor in body1 local frame; map to global frame + anchor1 = wp.vec3(data[0], data[1], data[2]) + pos = xpos_in[worldid, obj1id] + xmat_in[worldid, obj1id] @ anchor1 + + # data[3:6] = anchor position in body2 local frame + anchor2 = wp.transpose(xmat_in[worldid, obj2id]) @ (pos - xpos_in[worldid, obj2id]) + data[3] = anchor2[0] + data[4] = anchor2[1] + data[5] = anchor2[2] + eq_data_out[eq_data_id, eqid] = data + elif objtype == int(types.ObjType.SITE.value): + # site-based connect, eq_data is unused + eq_data_out[eq_data_id, eqid] = types.vec11(0.0) + elif eqtype == int(types.EqType.WELD.value): + if objtype == int(types.ObjType.BODY.value): + quat = wp.quat(data[6], data[7], data[8], data[9]) + if wp.length_sq(quat) > 0.0: + # user has set quaternion data: normalize it and keep the remaining data + quat = wp.normalize(quat) + data[6] = quat[0] + data[7] = quat[1] + data[8] = quat[2] + data[9] = quat[3] + eq_data_out[eq_data_id, eqid] = data + else: + obj1id = eq_obj1id[eqid] + obj2id = eq_obj2id[eqid] + + # data[0:3] = anchor in body2 local frame; map to global frame + anchor2 = wp.vec3(data[0], data[1], data[2]) + pos = xpos_in[worldid, obj2id] + xmat_in[worldid, obj2id] @ anchor2 + + # data[3:6] = anchor position in body1 local frame + anchor1 = wp.transpose(xmat_in[worldid, obj1id]) @ (pos - xpos_in[worldid, obj1id]) + data[3] = anchor1[0] + data[4] = anchor1[1] + data[5] = anchor1[2] + + # data[6:10] = neg(xquat1) * xquat2 = "xquat2 - xquat1" in body1 local frame + relquat = mjmath.mul_quat(mjmath.quat_inv(xquat_in[worldid, obj1id]), xquat_in[worldid, obj2id]) + data[6] = relquat[0] + data[7] = relquat[1] + data[8] = relquat[2] + data[9] = relquat[3] + eq_data_out[eq_data_id, eqid] = data + + @wp.kernel def _resolve_tendon_lengthspring( ten_length_in: wp.array2d[float], @@ -3005,6 +3164,8 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True): Computes: - tendon_length0: tendon resting lengths + - eq_data: connect/weld anchor data, recomputed so the constraint is + satisfied at qpos0 - dof_invweight0: inverse inertia for DOFs - body_invweight0: inverse spatial inertia for bodies - tendon_invweight0: inverse weight for tendons @@ -3044,6 +3205,13 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True): wp.launch(_copy_tendon_length0, dim=(d.nworld, m.ntendon), inputs=[d.ten_length], outputs=[m.tendon_length0]) + wp.launch( + _compute_eq_data0, + dim=(d.nworld, m.neq), + inputs=[m.eq_type, m.eq_obj1id, m.eq_obj2id, m.eq_objtype, d.xpos, d.xquat, d.xmat], + outputs=[m.eq_data], + ) + # dof_invweight0: computed per joint with averaging for multi-DOF joints # FREE: 6 DOFs, trans gets mean(A[0:3]), rot gets mean(A[3:6]) # BALL: 3 DOFs, all get mean(A[0:3]) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py index 059756ce..d5052f9b 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -652,6 +652,11 @@ def _flex_elasticity( elem_data_adr = flex_elemdataadr[f] + local_elemid * (dim + 1) vbase = flex_vertadr[f] + + # skip trilinear/interp elements (vertbodyid == -1, no simplex stiffness) + vert0_check = flex_elem[elem_data_adr] + if flex_vertbodyid[vbase + vert0_check] < 0: + return gradient = wp.matrix(0.0, shape=(6, 6)) for e in range(nedge): vert0 = flex_elem[elem_data_adr + edges[e, 0]] @@ -765,6 +770,180 @@ def _flex_bending( wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[i, x]) +@wp.kernel +def _flex_passive_interp( + # Model: + nflex: int, + body_rootid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + flex_interp: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], + flex_stiffnessadr: wp.array[int], + flex_nodebodyid: wp.array[int], + flex_node: wp.array[wp.vec3], + flex_node0: wp.array[wp.vec3], + flex_stiffness: wp.array[float], + flex_damping: wp.array[float], + flex_edgeequality: wp.array[int], + flex_centered: wp.array[bool], + flex_cell_map: wp.array[wp.vec4i], + # Data in: + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], + flexnode_xpos_in: wp.array2d[wp.vec3], + # In: + dsbl_spring: bool, + dsbl_damper: bool, + # Data out: + qfrc_spring_out: wp.array2d[float], + qfrc_damper_out: wp.array2d[float], + # Out: + displ_scratch_out: wp.array3d[wp.vec3], + vel_corot_scratch_out: wp.array3d[wp.vec3], +): + """Corotational passive forces for interpolated flex (trilinear/quadratic).""" + worldid, cellid = wp.tid() + + mapping = flex_cell_map[cellid] + f = mapping[0] + ci = mapping[1] + cj = mapping[2] + ck = mapping[3] + + order = flex_interp[f] + if order <= 0: + return + + npc = (order + 1) * (order + 1) * (order + 1) + ndof_cell = 3 * npc + + cellnum = flex_cellnum[f] + cy = cellnum[1] + cz = cellnum[2] + nstart = flex_nodeadr[f] + ny_g = cy * order + 1 + nz_g = cz * order + 1 + + # Cell stiffness matrix address + stiffness_adr_base = flex_stiffnessadr[f] + if stiffness_adr_base < 0: + return + + cell_idx = ci * cy * cz + cj * cz + ck + k_base = stiffness_adr_base + cell_idx * ndof_cell * ndof_cell + + # Skip empty cells (zero stiffness) + if flex_stiffness[k_base] == 0.0: + return + + cell_quat = support.compute_interp_cell_quat(flexnode_xpos_in, order, ci, cj, ck, cy, cz, ny_g, nz_g, nstart, worldid) + + # mju_negQuat: conjugate (R⁻¹) — negate xyz, keep w + cell_quat_inv = wp.quat(-cell_quat[0], -cell_quat[1], -cell_quat[2], cell_quat[3]) + + # Pre-compute displacements and velocities in corotational frame + # (matches C: rotate all positions/velocities once, then K*u) + idx_j = int(0) + for li_j in range(order + 1): + for lj_j in range(order + 1): + for lk_j in range(order + 1): + if idx_j < npc: + gi_j = ci * order + li_j + gj_j = cj * order + lj_j + gk_j = ck * order + lk_j + gidx_j = gi_j * ny_g * nz_g + gj_j * nz_g + gk_j + + xpos_j = flexnode_xpos_in[worldid, nstart + gidx_j] + + if not dsbl_spring: + refpos_j = flex_node0[nstart + gidx_j] + xrot_j = wp.quat_rotate(cell_quat_inv, xpos_j) + displ_scratch_out[worldid, cellid, idx_j] = xrot_j - refpos_j + + if not dsbl_damper: + bodyid_j = flex_nodebodyid[nstart + gidx_j] + cvel_j = cvel_in[worldid, bodyid_j] + omega_j = wp.spatial_top(cvel_j) + vcom_j = wp.spatial_bottom(cvel_j) + com_j = subtree_com_in[worldid, body_rootid[bodyid_j]] + r_j = xpos_j - com_j + vel_world_j = vcom_j + wp.cross(omega_j, r_j) + vel_corot_scratch_out[worldid, cellid, idx_j] = wp.quat_rotate(cell_quat_inv, vel_world_j) + + idx_j += 1 + + # Compute K*displacement and K*velocity per output node, then scatter forces + idx_i = int(0) + for li_i in range(order + 1): + for lj_i in range(order + 1): + for lk_i in range(order + 1): + if idx_i < npc: + gi_i = ci * order + li_i + gj_i = cj * order + lj_i + gk_i = ck * order + lk_i + gidx_i = gi_i * ny_g * nz_g + gj_i * nz_g + gk_i + bodyid_i = flex_nodebodyid[nstart + gidx_i] + + frc_spring = wp.vec3(0.0) + frc_damper = wp.vec3(0.0) + + for comp_i in range(3): + row = idx_i * 3 + comp_i + val_spring = float(0.0) + val_damper = float(0.0) + + for idx_j in range(npc): + for comp_j in range(3): + col = idx_j * 3 + comp_j + K_ij = flex_stiffness[k_base + row * ndof_cell + col] + + if not dsbl_spring: + val_spring += K_ij * displ_scratch_out[worldid, cellid, idx_j][comp_j] + + if not dsbl_damper: + val_damper += K_ij * vel_corot_scratch_out[worldid, cellid, idx_j][comp_j] + + frc_spring[comp_i] = val_spring + frc_damper[comp_i] = val_damper + + # Rotate forces back to world frame (R) + frc_spring_world = wp.quat_rotate(cell_quat, frc_spring) + frc_damper_world = wp.quat_rotate(cell_quat, frc_damper) + + # Scale damper force by damping coefficient + frc_damper_world = frc_damper_world * flex_damping[f] + + # Apply forces to body DOFs (fast path: nodes at body origin) + dofnum_i = body_dofnum[bodyid_i] + dofadr_i = body_dofadr[bodyid_i] + if dofnum_i > 0: + centered = flex_centered[f] + node_local = flex_node[nstart + gidx_i] + at_origin = node_local[0] == 0.0 and node_local[1] == 0.0 and node_local[2] == 0.0 + + if centered or at_origin: + for x in range(3): + if x < dofnum_i: + if not dsbl_spring: + wp.atomic_add( + qfrc_spring_out, + worldid, + dofadr_i + x, + frc_spring_world[x], + ) + if not dsbl_damper: + wp.atomic_add( + qfrc_damper_out, + worldid, + dofadr_i + x, + frc_damper_world[x], + ) + + idx_i += 1 + + @event_scope def passive(m: Model, d: Data): """Adds all passive forces.""" @@ -893,6 +1072,44 @@ def passive(m: Model, d: Data): outputs=[d.qfrc_gravcomp], ) + # Launch passive interp kernel for interpolated flex (trilinear/quadratic) + if m.nflex and m.nflexintcell > 0: + displ_scratch = wp.empty((d.nworld, m.nflexintcell, 27), dtype=wp.vec3) + vel_corot_scratch = wp.empty((d.nworld, m.nflexintcell, 27), dtype=wp.vec3) + wp.launch( + _flex_passive_interp, + dim=(d.nworld, m.nflexintcell), + inputs=[ + m.nflex, + m.body_rootid, + m.body_dofnum, + m.body_dofadr, + m.flex_interp, + m.flex_cellnum, + m.flex_nodeadr, + m.flex_stiffnessadr, + m.flex_nodebodyid, + m.flex_node, + m.flex_node0, + m.flex_stiffness, + m.flex_damping, + m.flex_edgeequality, + m.flex_centered, + m.flex_cell_map, + d.subtree_com, + d.cvel, + d.flexnode_xpos, + dsbl_spring, + dsbl_damper, + ], + outputs=[ + d.qfrc_spring, + d.qfrc_damper, + displ_scratch, + vel_corot_scratch, + ], + ) + if m.has_fluid: _fluid(m, d) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py index 4c2939e8..1defb28a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py @@ -228,14 +228,19 @@ def _site_local_to_global( def _flex_vertices( # Model: nflex: int, + flex_interp: wp.array[int], + flex_cellnum: wp.array[wp.vec3i], + flex_nodeadr: wp.array[int], flex_vertadr: wp.array[int], flex_vertnum: wp.array[int], flex_vertbodyid: wp.array[int], flex_vert: wp.array[wp.vec3], + flex_vert0: wp.array[wp.vec3], flex_centered: wp.array[bool], # Data in: xpos_in: wp.array2d[wp.vec3], xmat_in: wp.array2d[wp.mat33], + flexnode_xpos_in: wp.array2d[wp.vec3], # Data out: flexvert_xpos_out: wp.array2d[wp.vec3], ): @@ -246,15 +251,93 @@ def _flex_vertices( if locid >= 0 and locid < flex_vertnum[f]: break - bodyid = flex_vertbodyid[vertid] + if flex_interp[f] != 0: + # Interpolated flex: vertex position = weighted sum of node positions + coord = flex_vert0[vertid] + cn = flex_cellnum[f] + cx = cn[0] + cy = cn[1] + cz = cn[2] + + # Cell lookup: find containing cell + ci = wp.min(int(coord[0] * float(cx)), cx - 1) + ci = wp.max(ci, 0) + cj = wp.min(int(coord[1] * float(cy)), cy - 1) + cj = wp.max(cj, 0) + ck = wp.min(int(coord[2] * float(cz)), cz - 1) + ck = wp.max(ck, 0) + + # Local parametric coordinates within cell + local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0) + local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0) + local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0) + local = wp.vec3(local_x, local_y, local_z) + + # Node grid dimensions + ny_g = cy + 1 + nz_g = cz + 1 + nstart = flex_nodeadr[f] + + # Accumulate weighted node positions + result = wp.vec3(0.0, 0.0, 0.0) + for li in range(2): + for lj in range(2): + for lk in range(2): + w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk) + gi = ci + li + gj = cj + lj + gk = ck + lk + node_idx = gi * ny_g * nz_g + gj * nz_g + gk + result += w * flexnode_xpos_in[worldid, nstart + node_idx] + + flexvert_xpos_out[worldid, vertid] = result + else: + # Non-interpolated flex: vertex position from single body + bodyid = flex_vertbodyid[vertid] + xpos = xpos_in[worldid, bodyid] + + if flex_centered[f]: + flexvert_xpos_out[worldid, vertid] = xpos + else: + xmat = xmat_in[worldid, bodyid] + local_pos = flex_vert[vertid] + flexvert_xpos_out[worldid, vertid] = xmat @ local_pos + xpos + + +@wp.kernel +def _flex_nodes( + # Model: + nflex: int, + flex_nodeadr: wp.array[int], + flex_nodenum: wp.array[int], + flex_nodebodyid: wp.array[int], + flex_node: wp.array[wp.vec3], + flex_centered: wp.array[bool], + # Data in: + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + # Data out: + flexnode_xpos_out: wp.array2d[wp.vec3], +): + worldid, nodeid = wp.tid() + + for f in range(nflex): + locid = nodeid - flex_nodeadr[f] + if locid >= 0 and locid < flex_nodenum[f]: + break + + bodyid = flex_nodebodyid[nodeid] xpos = xpos_in[worldid, bodyid] if flex_centered[f]: - flexvert_xpos_out[worldid, vertid] = xpos + flexnode_xpos_out[worldid, nodeid] = xpos else: - xmat = xmat_in[worldid, bodyid] - local_pos = flex_vert[vertid] - flexvert_xpos_out[worldid, vertid] = xmat @ local_pos + xpos + local_pos = flex_node[nodeid] + if local_pos[0] == 0.0 and local_pos[1] == 0.0 and local_pos[2] == 0.0: + flexnode_xpos_out[worldid, nodeid] = xpos + else: + xmat = xmat_in[worldid, bodyid] + flexnode_xpos_out[worldid, nodeid] = xmat @ local_pos + xpos @wp.kernel @@ -302,6 +385,11 @@ def _flex_edges( b1 = flex_vertbodyid[vbase0] b2 = flex_vertbodyid[vbase1] + # skip Jacobian/velocity for trilinear flex (vertbodyid == -1) + if b1 < 0 or b2 < 0: + flexedge_velocity_out[worldid, edgeid] = 0.0 + return + dofnum1 = body_dofnum[b1] dofnum2 = body_dofnum[b2] @@ -417,18 +505,40 @@ def kinematics(m: Model, d: Data): @event_scope def flex(m: Model, d: Data): + # Compute node positions first (needed for interpolated vertex positions) + wp.launch( + _flex_nodes, + dim=(d.nworld, m.nflexnode), + inputs=[ + m.nflex, + m.flex_nodeadr, + m.flex_nodenum, + m.flex_nodebodyid, + m.flex_node, + m.flex_centered, + d.xpos, + d.xmat, + ], + outputs=[d.flexnode_xpos], + ) + wp.launch( _flex_vertices, dim=(d.nworld, m.nflexvert), inputs=[ m.nflex, + m.flex_interp, + m.flex_cellnum, + m.flex_nodeadr, m.flex_vertadr, m.flex_vertnum, m.flex_vertbodyid, m.flex_vert, + m.flex_vert0, m.flex_centered, d.xpos, d.xmat, + d.flexnode_xpos, ], outputs=[d.flexvert_xpos], ) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py index f6ae3f7a..d66ca3c9 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -87,6 +87,7 @@ def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext: jv=wp.empty((nworld, njmax), dtype=float), quad=wp.empty((nworld, njmax), dtype=wp.vec3), alpha=wp.empty((nworld,), dtype=float), + grad_scale=wp.empty((nworld,), dtype=float), improvement=wp.empty((nworld,), dtype=float), prev_grad=wp.empty((nworld, nv), dtype=float), prev_Mgrad=wp.empty((nworld, nv), dtype=float), @@ -839,6 +840,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType, ctx_jv_out: wp.array2d[float], ctx_quad_out: wp.array2d[wp.vec3], ctx_improvement_out: wp.array[float], + ctx_alpha_out: wp.array[float], ): worldid, tid = wp.tid() @@ -1202,7 +1204,11 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType, swap_hi = swap_hi_hi_next or swap_hi_mid or swap_hi_lo_next # check for convergence - ls_done = (not swap_lo and not swap_hi) or (lo[1] < 0.0 and lo[1] > -gtol) or (hi[1] > 0.0 and hi[1] < gtol) + ls_done = ( + (not swap_lo and not swap_hi) + or (lo[0] < 0.0 and lo[1] < 0.0 and lo[1] > -gtol) + or (hi[0] < 0.0 and hi[1] > 0.0 and hi[1] < gtol) + ) # update alpha if improved improved = lo[0] < 0.0 or hi[0] < 0.0 @@ -1229,6 +1235,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType, if tid == 0: ctx_improvement_out[worldid] = improvement + ctx_alpha_out[worldid] = alpha return kernel @@ -1276,7 +1283,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus ctx.quad, ctx.done, ], - outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement], + outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement, ctx.alpha], block_dim=m.block_dim.linesearch_iterative, ) @@ -1380,8 +1387,13 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext): # jv = J @ search (when not fused into iterative kernel) if not fuse_jv: - dofs_per_thread = 20 if m.nv > 50 else 50 - threads_per_efc = ceil(m.nv / dofs_per_thread) + if m.is_sparse: + # Sparse J has few nonzeros per row, one thread handles them all. + dofs_per_thread = m.nv + threads_per_efc = 1 + else: + dofs_per_thread = 20 if m.nv > 50 else 50 + threads_per_efc = ceil(m.nv / dofs_per_thread) if threads_per_efc > 1: wp.launch( @@ -1636,6 +1648,7 @@ def _update_constraint_init_qfrc_constraint_sparse( efc_J_in: wp.array3d[float], efc_force_in: wp.array2d[float], # In: + changed_count_in: wp.array[int], ctx_done_in: wp.array[bool], # Data out: qfrc_constraint_out: wp.array2d[float], @@ -1645,10 +1658,15 @@ def _update_constraint_init_qfrc_constraint_sparse( if ctx_done_in[worldid]: return + if changed_count_in[worldid] == 0: + return + if efcid >= nefc_in[worldid]: return force = efc_force_in[worldid, efcid] + if force == 0.0: + return rownnz = efc_J_rownnz_in[worldid, efcid] rowadr = efc_J_rowadr_in[worldid, efcid] @@ -1660,29 +1678,58 @@ def _update_constraint_init_qfrc_constraint_sparse( @wp.kernel -def _update_constraint_init_qfrc_constraint_dense( +def _qfrc_constraint_from_grad( # Data in: - nefc_in: wp.array[int], - efc_J_in: wp.array3d[float], - efc_force_in: wp.array2d[float], - njmax_in: int, + qfrc_smooth_in: wp.array2d[float], + efc_Ma_in: wp.array2d[float], # In: - ctx_done_in: wp.array[bool], + ctx_grad_in: wp.array2d[float], + ctx_grad_scale_in: wp.array[float], # Data out: qfrc_constraint_out: wp.array2d[float], ): worldid, dofid = wp.tid() - if ctx_done_in[worldid]: - return + grad = ctx_grad_scale_in[worldid] * ctx_grad_in[worldid, dofid] + qfrc_constraint_out[worldid, dofid] = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - grad - sum_qfrc = float(0.0) - for efcid in range(min(njmax_in, nefc_in[worldid])): - efc_J = efc_J_in[worldid, efcid, dofid] - force = efc_force_in[worldid, efcid] - sum_qfrc += efc_J * force - qfrc_constraint_out[worldid, dofid] = sum_qfrc +@cache_kernel +def _update_constraint_init_qfrc_constraint_dense(stable_fast: bool): + STABLE_FAST = stable_fast + + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # Data in: + nefc_in: wp.array[int], + efc_J_in: wp.array3d[float], + efc_force_in: wp.array2d[float], + njmax_in: int, + # In: + changed_count_in: wp.array[int], + ctx_done_in: wp.array[bool], + # Data out: + qfrc_constraint_out: wp.array2d[float], + ): + worldid, dofid = wp.tid() + + if ctx_done_in[worldid]: + return + + # Fast path: stale qfrc_constraint is never read; recovered after the solve. + if wp.static(STABLE_FAST): + if changed_count_in[worldid] == 0: + return + + sum_qfrc = float(0.0) + for efcid in range(min(njmax_in, nefc_in[worldid])): + efc_J = efc_J_in[worldid, efcid, dofid] + force = efc_force_in[worldid, efcid] + sum_qfrc += efc_J * force + + qfrc_constraint_out[worldid, dofid] = sum_qfrc + + return kernel @wp.kernel @@ -1744,48 +1791,51 @@ def _update_gradient_h_incremental_sparse( # In: changed_ids_in: wp.array2d[int], changed_count_in: wp.array[int], + slots_per_world: int, # Out: ctx_h_out: wp.array3d[float], ): - """Incrementally update upper triangle of H for changed constraints (sparse J).""" - worldid, change_idx = wp.tid() + """Incrementally update upper triangle of H for changed constraints (sparse J). + + One warp per changed constraint row: the lanes split the row's upper-triangular + entries (same sqrt triangular-number decode as _JTDAJ_sparse), replacing the + serial nnz^2 loop that dominated this kernel. + """ + worldid, slot, lane = wp.tid() n_changes = changed_count_in[worldid] - if change_idx >= n_changes: - return + for change_idx in range(slot, n_changes, slots_per_world): + efcid = changed_ids_in[worldid, change_idx] + D = efc_D_in[worldid, efcid] + sign = float(0.0) + if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value: + sign = D + else: + sign = -D - efcid = changed_ids_in[worldid, change_idx] - D = efc_D_in[worldid, efcid] - sign = float(0.0) - if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value: - sign = D - else: - sign = -D + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + n_entries = rownnz * (rownnz + 1) // 2 - rownnz = efc_J_rownnz_in[worldid, efcid] - rowadr = efc_J_rowadr_in[worldid, efcid] - - for ii in range(rownnz): - sparseidi = rowadr + ii - Ji = efc_J_in[worldid, 0, sparseidi] - if Ji == 0.0: - continue - colindi = efc_J_colind_in[worldid, 0, sparseidi] - for jj in range(ii + 1): - sparseidj = rowadr + jj - Jj = efc_J_in[worldid, 0, sparseidj] - if Jj == 0.0: - continue - colindj = efc_J_colind_in[worldid, 0, sparseidj] + for entry in range(lane, n_entries, wp.static(_JTDAJ_THREADS_PER_GROUP)): + ii = int((wp.sqrt(float(8 * entry + 1)) - 1.0) * 0.5) + jj = entry - ii * (ii + 1) // 2 + Ji = efc_J_in[worldid, 0, rowadr + ii] + Jj = efc_J_in[worldid, 0, rowadr + jj] h = sign * Ji * Jj - # Ensure upper triangle: smaller index first. - if colindi <= colindj: - wp.atomic_add(ctx_h_out[worldid, colindi], colindj, h) - else: - wp.atomic_add(ctx_h_out[worldid, colindj], colindi, h) + if h != 0.0: + colindi = efc_J_colind_in[worldid, 0, rowadr + ii] + colindj = efc_J_colind_in[worldid, 0, rowadr + jj] + wp.atomic_add(ctx_h_out[worldid, wp.min(colindi, colindj)], wp.max(colindi, colindj), h) -def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, track_changes: bool = False): +def _update_constraint( + m: types.Model, + d: types.Data, + ctx: SolverContext | InverseContext, + track_changes: bool = False, + stable_fast: bool = False, +): """Update constraint arrays after each solve iteration.""" efc_inputs = [ m.opt.impratio_invsqrt, @@ -1811,59 +1861,97 @@ def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | Inver outputs=[d.efc.force, d.efc.state, ctx.changed_efc_ids, ctx.changed_efc_count], ) - # qfrc_constraint = efc_J.T @ efc_force + # qfrc_constraint = efc_J.T @ efc_force. Fast-path worlds with no state flips + # skip the rebuild; the public value is recovered after the solve. + changed = ctx.changed_efc_count if stable_fast else d.nefc if m.is_sparse: d.qfrc_constraint.zero_() wp.launch( _update_constraint_init_qfrc_constraint_sparse, dim=(d.nworld, d.njmax), - inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, ctx.done], + inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, changed, ctx.done], outputs=[d.qfrc_constraint], ) else: wp.launch( - _update_constraint_init_qfrc_constraint_dense, + _update_constraint_init_qfrc_constraint_dense(stable_fast), dim=(d.nworld, m.nv), - inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, ctx.done], + inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, changed, ctx.done], outputs=[d.qfrc_constraint], ) -@wp.kernel -def _update_gradient_zero_grad_dot( - # In: - ctx_done_in: wp.array[bool], - # Out: - ctx_grad_dot_out: wp.array[float], -): - worldid = wp.tid() +@cache_kernel +def _update_gradient_zero_grad_dot(stable_fast: bool): + STABLE_FAST = stable_fast - if ctx_done_in[worldid]: - return + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # In: + changed_count_in: wp.array[int], + ctx_alpha_in: wp.array[float], + ctx_done_in: wp.array[bool], + # Out: + ctx_grad_dot_out: wp.array[float], + ctx_grad_scale_out: wp.array[float], + ): + worldid = wp.tid() - ctx_grad_dot_out[worldid] = 0.0 + if ctx_done_in[worldid]: + return + + # Fast path: grad stays stale at its last rebuilt value g. The true + # gradient is grad_scale * g, and a linesearch step t along the (equally + # stale) search direction changes it to (grad_scale - t) * g. + if wp.static(STABLE_FAST): + if changed_count_in[worldid] == 0: + sigma = ctx_grad_scale_out[worldid] + new_sigma = sigma - ctx_alpha_in[worldid] + ratio = float(0.0) + if sigma != 0.0: + ratio = new_sigma / sigma + ctx_grad_dot_out[worldid] *= ratio * ratio + ctx_grad_scale_out[worldid] = new_sigma + return + + ctx_grad_dot_out[worldid] = 0.0 + ctx_grad_scale_out[worldid] = 1.0 + + return kernel -@wp.kernel -def _update_gradient_grad( - # Data in: - qfrc_smooth_in: wp.array2d[float], - qfrc_constraint_in: wp.array2d[float], - efc_Ma_in: wp.array2d[float], - # In: - ctx_done_in: wp.array[bool], - # Out: - ctx_grad_out: wp.array2d[float], - ctx_grad_dot_out: wp.array[float], -): - worldid, dofid = wp.tid() +@cache_kernel +def _update_gradient_grad(stable_fast: bool): + STABLE_FAST = stable_fast - if ctx_done_in[worldid]: - return + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # Data in: + qfrc_smooth_in: wp.array2d[float], + qfrc_constraint_in: wp.array2d[float], + efc_Ma_in: wp.array2d[float], + # In: + changed_count_in: wp.array[int], + ctx_done_in: wp.array[bool], + # Out: + ctx_grad_out: wp.array2d[float], + ctx_grad_dot_out: wp.array[float], + ): + worldid, dofid = wp.tid() - grad = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - qfrc_constraint_in[worldid, dofid] - ctx_grad_out[worldid, dofid] = grad - wp.atomic_add(ctx_grad_dot_out, worldid, grad * grad) + if ctx_done_in[worldid]: + return + + # Fast path: grad stays stale (see _update_gradient_zero_grad_dot). + if wp.static(STABLE_FAST): + if changed_count_in[worldid] == 0: + return + + grad = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - qfrc_constraint_in[worldid, dofid] + ctx_grad_out[worldid, dofid] = grad + wp.atomic_add(ctx_grad_dot_out, worldid, grad * grad) + + return kernel @wp.kernel @@ -2585,12 +2673,15 @@ def _update_gradient_JTCJ_dense( @cache_kernel -def _update_gradient_cholesky(tile_size: int): +def _update_gradient_cholesky(tile_size: int, skip_noflip: bool = False): + SKIP_NOFLIP = skip_noflip + @wp.kernel(module="unique", enable_backward=False) def kernel( # In: ctx_grad_in: wp.array2d[float], h_in: wp.array3d[float], + changed_count_in: wp.array[int], ctx_done_in: wp.array[bool], # Out: ctx_Mgrad_out: wp.array2d[float], @@ -2601,6 +2692,11 @@ def _update_gradient_cholesky(tile_size: int): if ctx_done_in[worldid]: return + # Fast path: skip the solve (see the blocked skip_unchanged variant). + if wp.static(SKIP_NOFLIP): + if changed_count_in[worldid] == 0: + return + mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE)) wp.tile_cholesky_inplace(mat_tile, fill_mode="upper") input_tile = wp.tile_load(ctx_grad_in[worldid], shape=TILE_SIZE) @@ -2642,8 +2738,9 @@ def _update_gradient_cholesky_blocked(tile_size: int, matrix_size: int, check_sk @cache_kernel -def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int): +def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int, skip_noflip: bool = False): """Blocked Cholesky that skips factorization when no constraints changed.""" + SKIP_NOFLIP = skip_noflip @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False}) def kernel( @@ -2662,14 +2759,24 @@ def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size if ctx_done_in[worldid]: return - if changed_count_in[worldid] > 0: + # Fast path: skip the solve; Mgrad stays stale on the unchanged ray, and + # the linesearch is invariant to the scale of its direction. + if wp.static(SKIP_NOFLIP): + if changed_count_in[worldid] == 0: + return + wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))( ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid] ) else: - wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))( - ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid] - ) + if changed_count_in[worldid] > 0: + wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))( + ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid] + ) + else: + wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))( + ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid] + ) return kernel @@ -2685,7 +2792,9 @@ def _padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float ctx_h_out[worldid, dofid, dofid] = 1.0 -def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext, skip_unchanged: bool = False): +def _cholesky_factorize_solve( + m: types.Model, d: types.Data, ctx: SolverContext, skip_unchanged: bool = False, skip_noflip: bool = False +): """Cholesky factorize ctx.h and solve for Mgrad. If skip_unchanged is True (blocked path only), worlds where no constraints @@ -2693,9 +2802,9 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext, """ if m.nv <= _BLOCK_CHOLESKY_DIM: wp.launch_tiled( - _update_gradient_cholesky(m.nv), + _update_gradient_cholesky(m.nv, skip_noflip), dim=d.nworld, - inputs=[ctx.grad, ctx.h, ctx.done], + inputs=[ctx.grad, ctx.h, ctx.changed_efc_count if skip_noflip else d.nefc, ctx.done], outputs=[ctx.Mgrad], block_dim=m.block_dim.update_gradient_cholesky, ) @@ -2709,7 +2818,7 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext, if skip_unchanged: wp.launch_tiled( - _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad), + _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad, skip_noflip), dim=d.nworld, inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.changed_efc_count, ctx.hfactor], outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))], @@ -2806,11 +2915,16 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact: block_dim=m.block_dim.update_gradient_grad, ) else: - wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot]) wp.launch( - _update_gradient_grad, + _update_gradient_zero_grad_dot(False), + dim=d.nworld, + inputs=[d.nefc, ctx.alpha, ctx.done], + outputs=[ctx.grad_dot, ctx.grad_scale], + ) + wp.launch( + _update_gradient_grad(False), dim=(d.nworld, m.nv), - inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done], + inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, d.nefc, ctx.done], outputs=[ctx.grad, ctx.grad_dot], ) @@ -3002,26 +3116,33 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact: raise ValueError(f"Unknown solver type: {m.opt.solver}") -def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverContext): +def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverContext, stable_fast: bool = False): """Incremental gradient update: update H for changed constraints + re-factorize. Skips the full J^T*D*J rebuild by applying only the delta from constraints that changed QUADRATIC state, then re-factorizes and solves. """ - wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot]) + changed = ctx.changed_efc_count if stable_fast else d.nefc + wp.launch( + _update_gradient_zero_grad_dot(stable_fast), + dim=d.nworld, + inputs=[changed, ctx.alpha, ctx.done], + outputs=[ctx.grad_dot, ctx.grad_scale], + ) wp.launch( - _update_gradient_grad, + _update_gradient_grad(stable_fast), dim=(d.nworld, m.nv), - inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done], + inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, changed, ctx.done], outputs=[ctx.grad, ctx.grad_dot], ) # Update upper triangle of H with delta from changed constraints. if m.is_sparse: + slots = _jtdaj_groups_per_world(d.nworld, ctx.changed_efc_ids.shape[1]) wp.launch( _update_gradient_h_incremental_sparse, - dim=(d.nworld, ctx.changed_efc_ids.shape[1]), + dim=(d.nworld, slots, _JTDAJ_THREADS_PER_GROUP), inputs=[ d.efc.J_rownnz, d.efc.J_rowadr, @@ -3031,6 +3152,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte d.efc.state, ctx.changed_efc_ids, ctx.changed_efc_count, + slots, ], outputs=[ctx.h], ) @@ -3049,7 +3171,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte outputs=[ctx.h], ) - _cholesky_factorize_solve(m, d, ctx, skip_unchanged=True) + _cholesky_factorize_solve(m, d, ctx, skip_unchanged=True, skip_noflip=stable_fast) @wp.kernel @@ -3128,46 +3250,70 @@ def _solve_beta_accumulate( wp.atomic_add(ctx_beta_den_out, worldid, den) -@wp.kernel -def _solve_zero_search_dot( - # In: - ctx_done_in: wp.array[bool], - # Out: - ctx_search_dot_out: wp.array[float], -): - worldid = wp.tid() +@cache_kernel +def _solve_zero_search_dot(stable_fast: bool): + STABLE_FAST = stable_fast - if ctx_done_in[worldid]: - return + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # In: + changed_count_in: wp.array[int], + ctx_done_in: wp.array[bool], + # Out: + ctx_search_dot_out: wp.array[float], + ): + worldid = wp.tid() - ctx_search_dot_out[worldid] = 0.0 + if ctx_done_in[worldid]: + return + + # Fast path: search stays on the same ray; keep search_dot consistent with it. + if wp.static(STABLE_FAST): + if changed_count_in[worldid] == 0: + return + + ctx_search_dot_out[worldid] = 0.0 + + return kernel -@wp.kernel -def _solve_search_update( - # Model: - opt_solver: int, - # In: - ctx_Mgrad_in: wp.array2d[float], - ctx_search_in: wp.array2d[float], - ctx_beta_in: wp.array[float], - ctx_done_in: wp.array[bool], - # Out: - ctx_search_out: wp.array2d[float], - ctx_search_dot_out: wp.array[float], -): - worldid, dofid = wp.tid() +@cache_kernel +def _solve_search_update(stable_fast: bool): + STABLE_FAST = stable_fast - if ctx_done_in[worldid]: - return + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # Model: + opt_solver: int, + # In: + changed_count_in: wp.array[int], + ctx_Mgrad_in: wp.array2d[float], + ctx_search_in: wp.array2d[float], + ctx_beta_in: wp.array[float], + ctx_done_in: wp.array[bool], + # Out: + ctx_search_out: wp.array2d[float], + ctx_search_dot_out: wp.array[float], + ): + worldid, dofid = wp.tid() - search = -1.0 * ctx_Mgrad_in[worldid, dofid] + if ctx_done_in[worldid]: + return - if opt_solver == types.SolverType.CG: - search += ctx_beta_in[worldid] * ctx_search_in[worldid, dofid] + # Fast path: search stays on the stale ray; the linesearch absorbs its scale. + if wp.static(STABLE_FAST): + if changed_count_in[worldid] == 0: + return - ctx_search_out[worldid, dofid] = search - wp.atomic_add(ctx_search_dot_out, worldid, search * search) + search = -1.0 * ctx_Mgrad_in[worldid, dofid] + + if opt_solver == types.SolverType.CG: + search += ctx_beta_in[worldid] * ctx_search_in[worldid, dofid] + + ctx_search_out[worldid, dofid] = search + wp.atomic_add(ctx_search_dot_out, worldid, search * search) + + return kernel @wp.kernel @@ -3291,6 +3437,16 @@ def _solve_done( wp.atomic_add(nsolving_out, 0, -1) +def _use_incremental(m: types.Model) -> bool: + """Whether constraint state changes are tracked for incremental H updates.""" + return m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC + + +def _stable_fast(m: types.Model, compact: bool) -> bool: + """Stable-state fast path: needs state-change tracking; compact scatters qfrc itself.""" + return _use_incremental(m) and not compact + + @event_scope def _solver_iteration( m: types.Model, @@ -3305,16 +3461,21 @@ def _solver_iteration( # path in _update_constraint_efc has early returns that skip state change # tracking, and the additional JTCJ Hessian term depends on Jaref which # changes every iteration. - incremental = m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC + incremental = _use_incremental(m) + # Stable-state fast path: worlds with no state flips this iteration were + # exactly quadratic over the step, so grad/Mgrad/search only changed by a + # scalar along the same ray. Skip their qfrc/grad/solve/search updates and + # track the scalar in ctx.grad_scale. + stable_fast = _stable_fast(m, compact) if incremental: # Must complete before _update_constraint_efc which atomically increments. ctx.changed_efc_count.zero_() - _update_constraint(m, d, ctx, track_changes=incremental) + _update_constraint(m, d, ctx, track_changes=incremental, stable_fast=stable_fast) if incremental: - _update_gradient_incremental(m, d, ctx) + _update_gradient_incremental(m, d, ctx, stable_fast) else: _update_gradient(m, d, ctx, compact=compact) @@ -3362,12 +3523,13 @@ def _solver_iteration( ) else: - wp.launch(_solve_zero_search_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.search_dot]) + changed = ctx.changed_efc_count if stable_fast else d.nefc + wp.launch(_solve_zero_search_dot(stable_fast), dim=d.nworld, inputs=[changed, ctx.done], outputs=[ctx.search_dot]) wp.launch( - _solve_search_update, + _solve_search_update(stable_fast), dim=(d.nworld, m.nv), - inputs=[m.opt.solver, ctx.Mgrad, ctx.search, ctx.beta, ctx.done], + inputs=[m.opt.solver, changed, ctx.Mgrad, ctx.search, ctx.beta, ctx.done], outputs=[ctx.search, ctx.search_dot], ) @@ -3400,12 +3562,16 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont # if we are only using 1 thread, it makes sense to do more dofs as we can also skip the # init kernel. For more than 1 thread, dofs_per_thread is lower for better load balancing. - if m.nv > 50: + if m.is_sparse: + # Sparse J has few nonzeros per row, one thread handles them all. + dofs_per_thread = m.nv + threads_per_efc = 1 + elif m.nv > 50: dofs_per_thread = 20 + threads_per_efc = ceil(m.nv / dofs_per_thread) else: dofs_per_thread = 50 - - threads_per_efc = ceil(m.nv / dofs_per_thread) + threads_per_efc = ceil(m.nv / dofs_per_thread) # we need to clear the jaref array if we're doing atomic adds. if threads_per_efc > 1: ctx.Jaref.zero_() @@ -3490,6 +3656,16 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext, compact: bool = Fa for _ in range(m.opt.iterations): _solver_iteration(m, d, ctx, nsolving, compact=compact) + # Recover the public qfrc_constraint: the fast path leaves it stale, and the + # per-iteration zeroing wiped it for worlds that converged early. + if _stable_fast(m, compact): + wp.launch( + _qfrc_constraint_from_grad, + dim=(d.nworld, m.nv), + inputs=[d.qfrc_smooth, d.efc.Ma, ctx.grad, ctx.grad_scale], + outputs=[d.qfrc_constraint], + ) + # Active-DOF compaction solve (nvmax < nv). # diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py index ab7d5c62..4dc91e51 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -64,6 +64,115 @@ def next_act( return act +@wp.func +def mat33_to_quat_polar(F: wp.mat33) -> wp.quat: + cell_quat = wp.quat(0.0, 0.0, 0.0, 1.0) + for _iter in range(10): + rot = wp.quat_to_matrix(cell_quat) + rot_t = wp.transpose(rot) + col1_rot = rot_t[0] + col2_rot = rot_t[1] + col3_rot = rot_t[2] + F_t = wp.transpose(F) + col1_mat = F_t[0] + col2_mat = F_t[1] + col3_mat = F_t[2] + + omega = wp.cross(col1_rot, col1_mat) + wp.cross(col2_rot, col2_mat) + wp.cross(col3_rot, col3_mat) + denom = wp.abs(wp.dot(col1_rot, col1_mat) + wp.dot(col2_rot, col2_mat) + wp.dot(col3_rot, col3_mat)) + 1.0e-10 + omega = omega / denom + + w = wp.length(omega) + if w < 1.0e-6: + break + + axis = omega / w + half_w = 0.5 * w + qrot = wp.quat( + axis[0] * wp.sin(half_w), + axis[1] * wp.sin(half_w), + axis[2] * wp.sin(half_w), + wp.cos(half_w), + ) + cell_quat = wp.normalize(qrot * cell_quat) + return cell_quat + + +@wp.func +def compute_interp_cell_quat( + # Data in: + flexnode_xpos_in: wp.array2d[wp.vec3], + # In: + order: int, + ci: int, + cj: int, + ck: int, + cy: int, + cz: int, + ny_g: int, + nz_g: int, + nstart: int, + worldid: int, +) -> wp.quat: + """Computes corotational cell quaternion from deformation gradient at cell center.""" + npc = (order + 1) * (order + 1) * (order + 1) + F = wp.mat33(0.0) + idx = int(0) + for li in range(order + 1): + for lj in range(order + 1): + for lk in range(order + 1): + if idx < npc: + gi = ci * order + li + gj = cj * order + lj + gk = ck * order + lk + gidx = gi * ny_g * nz_g + gj * nz_g + gk + + node_pos = flexnode_xpos_in[worldid, nstart + gidx] + + if order == 1: + dphi_x = float(-1) if li == 0 else float(1) + dphi_y = float(-1) if lj == 0 else float(1) + dphi_z = float(-1) if lk == 0 else float(1) + phi_x = float(0.5) + phi_y = float(0.5) + phi_z = float(0.5) + else: + if li == 0: + dphi_x = -1.0 + elif li == 1: + dphi_x = 0.0 + else: + dphi_x = 1.0 + if lj == 0: + dphi_y = -1.0 + elif lj == 1: + dphi_y = 0.0 + else: + dphi_y = 1.0 + if lk == 0: + dphi_z = -1.0 + elif lk == 1: + dphi_z = 0.0 + else: + dphi_z = 1.0 + phi_x = 0.5 if li == 0 or li == 2 else 1.0 + phi_y = 0.5 if lj == 0 or lj == 2 else 1.0 + phi_z = 0.5 if lk == 0 or lk == 2 else 1.0 + + grad_x = dphi_x * phi_y * phi_z + grad_y = phi_x * dphi_y * phi_z + grad_z = phi_x * phi_y * dphi_z + + for r in range(3): + F[r, 0] += node_pos[r] * grad_x + F[r, 1] += node_pos[r] * grad_y + F[r, 2] += node_pos[r] * grad_z + + idx += 1 + + return mat33_to_quat_polar(F) + + @cache_kernel def mul_m_kernel(check_skip: bool): @wp.kernel(module="unique") @@ -893,3 +1002,84 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt d.userdata, ], ) + + +@wp.func +def _phi(s: float, i: int) -> float: + """1D trilinear basis function (order=1 only). + + phi(s, 0) = 1 - s + phi(s, 1) = s + """ + if i == 0: + return 1.0 - s + return s + + +@wp.func +def eval_basis_trilinear(local: wp.vec3, node_idx: int) -> float: + """Evaluate trilinear basis function for node_idx at local coords [0,1]^3. + + For order=1 (trilinear), node_idx encodes (i,j,k) via bits: + k = node_idx & 1, j = (node_idx >> 1) & 1, i = (node_idx >> 2) & 1 + """ + k = node_idx & 1 + j = (node_idx >> 1) & 1 + i = (node_idx >> 2) & 1 + return _phi(local[0], i) * _phi(local[1], j) * _phi(local[2], k) + + +@wp.func +def select_top4_weights( + # In: + W_mat: wp.mat33, + b_mat: wp.mat33, +) -> tuple[wp.vec4i, wp.vec4]: + """Selects top 4 weights and their corresponding body IDs from 8 voxel corners.""" + selected_b = wp.vec4i(-1, -1, -1, -1) + selected_W = wp.vec4(0.0, 0.0, 0.0, 0.0) + + local_W = W_mat + for p in range(4): + max_w = -1.0 + max_b = -1 + max_r = -1 + max_c = -1 + for r in range(3): + for c in range(3): + idx = 3 * r + c + if idx < 8: + w = local_W[r, c] + if w > max_w: + max_w = w + max_b = int(b_mat[r, c]) + max_r = r + max_c = c + # Record top choice for this pass and mark it as visited + if max_r >= 0: + local_W[max_r, max_c] = -1.0 + + if p == 0: + selected_b = wp.vec4i(max_b, -1, -1, -1) + selected_W = wp.vec4(max_w, 0.0, 0.0, 0.0) + elif p == 1: + selected_b = wp.vec4i(selected_b[0], max_b, -1, -1) + selected_W = wp.vec4(selected_W[0], max_w, 0.0, 0.0) + elif p == 2: + selected_b = wp.vec4i(selected_b[0], selected_b[1], max_b, -1) + selected_W = wp.vec4(selected_W[0], selected_W[1], max_w, 0.0) + else: + selected_b = wp.vec4i(selected_b[0], selected_b[1], selected_b[2], max_b) + selected_W = wp.vec4(selected_W[0], selected_W[1], selected_W[2], max_w) + + # Normalize selected weights + sum_W = selected_W[0] + selected_W[1] + selected_W[2] + selected_W[3] + if sum_W > 1.0e-5: + selected_W = wp.vec4( + selected_W[0] / sum_W, + selected_W[1] / sum_W, + selected_W[2] / sum_W, + selected_W[3] / sum_W, + ) + + return selected_b, selected_W diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py new file mode 100644 index 00000000..55a6a732 --- /dev/null +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py @@ -0,0 +1,351 @@ +# Copyright 2025 The Newton Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Tests for flex.""" + +import mujoco +import numpy as np +from absl.testing import absltest +from absl.testing import parameterized + +import mujoco.mjx.third_party.mujoco_warp as mjw +from mujoco.mjx.third_party.mujoco_warp import test_data + +_TRILINEAR_STRAIN_XML = """ + + + + + + + + + +""" + +# tolerance for difference between MuJoCo and MJWarp, mostly due to float32 +_TOLERANCE = 5e-4 + + +class TrilinearFlexTest(parameterized.TestCase): + def test_flexstrain_constraint_at_rest(self): + """Test FLEXSTRAIN constraint count and residuals match MuJoCo at rest.""" + mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML) + mjd = mujoco.MjData(mjm) + mujoco.mj_forward(mjm, mjd) + + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd) + + mjw.fwd_position(m, d) + mjw.make_constraint(m, d) + + # constraint counts should match + ne_warp = d.ne.numpy()[0] + self.assertEqual(ne_warp, mjd.ne, f"ne mismatch: warp={ne_warp}, mj={mjd.ne}") + + nefc_warp = d.nefc.numpy()[0] + self.assertEqual(nefc_warp, mjd.nefc, f"nefc mismatch: warp={nefc_warp}, mj={mjd.nefc}") + + # residuals should match + efc_pos = d.efc.pos.numpy()[0, :nefc_warp] + efc_pos_mj = mjd.efc_pos[: mjd.nefc] + np.testing.assert_allclose(efc_pos, efc_pos_mj, atol=1e-5, err_msg="FLEXSTRAIN residuals should match MuJoCo at rest") + + def test_flexstrain_constraint_perturbed(self): + """Test FLEXSTRAIN residuals and Jacobians match MuJoCo under perturbation.""" + mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML) + mjd = mujoco.MjData(mjm) + + # perturb first node + mjd.qpos[0] += 0.01 + mjd.qpos[3] += 0.005 + mujoco.mj_forward(mjm, mjd) + + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd) + + mjw.fwd_position(m, d) + mjw.make_constraint(m, d) + + nefc = d.nefc.numpy()[0] + nv = mjm.nv + self.assertEqual(nefc, mjd.nefc) + + # residuals + efc_pos_warp = d.efc.pos.numpy()[0, :nefc] + efc_pos_mj = mjd.efc_pos[: mjd.nefc] + np.testing.assert_allclose( + efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under perturbation" + ) + + # Jacobians + if mujoco.mj_isSparse(mjm): + mj_efc_J = np.zeros((mjd.nefc, nv)) + mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind) + else: + mj_efc_J = mjd.efc_J.reshape((mjd.nefc, nv)) + + if m.is_sparse: + warp_efc_J = np.zeros((nefc, nv)) + mujoco.mju_sparse2dense( + warp_efc_J, + d.efc.J.numpy()[0, 0], + d.efc.J_rownnz.numpy()[0, :nefc], + d.efc.J_rowadr.numpy()[0, :nefc], + d.efc.J_colind.numpy()[0, 0], + ) + else: + warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv] + + np.testing.assert_allclose(warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo") + + def test_flexstrain_constraint_rotated(self): + """Test FLEXSTRAIN residuals and Jacobians match MuJoCo under large rotation perturbation.""" + mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML) + mjd = mujoco.MjData(mjm) + + # Apply a rotation perturbation: rotate all node positions around Y axis by 30 degrees + # (0.5235 radians) + theta = 0.5235 + cos_t = np.cos(theta) + sin_t = np.sin(theta) + for i in range(0, mjm.nq, 3): + x = mjd.qpos[i] + z = mjd.qpos[i + 2] + mjd.qpos[i] = x * cos_t - z * sin_t + mjd.qpos[i + 2] = x * sin_t + z * cos_t + + mujoco.mj_forward(mjm, mjd) + + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd) + + mjw.fwd_position(m, d) + mjw.make_constraint(m, d) + + nefc = d.nefc.numpy()[0] + nv = mjm.nv + self.assertEqual(nefc, mjd.nefc) + + # residuals + efc_pos_warp = d.efc.pos.numpy()[0, :nefc] + efc_pos_mj = mjd.efc_pos[: mjd.nefc] + np.testing.assert_allclose( + efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under rotation" + ) + + # Jacobians + if mujoco.mj_isSparse(mjm): + mj_efc_J = np.zeros((mjd.nefc, nv)) + mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind) + else: + mj_efc_J = mjd.efc_J.reshape((mjd.nefc, nv)) + + if m.is_sparse: + warp_efc_J = np.zeros((nefc, nv)) + mujoco.mju_sparse2dense( + warp_efc_J, + d.efc.J.numpy()[0, 0], + d.efc.J_rownnz.numpy()[0, :nefc], + d.efc.J_rowadr.numpy()[0, :nefc], + d.efc.J_colind.numpy()[0, 0], + ) + else: + warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv] + + np.testing.assert_allclose( + warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo under rotation" + ) + + def test_flexstrain_rotational_invariance(self): + """Test that FLEXSTRAIN residuals are invariant under rigid translation.""" + mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML) + mjd = mujoco.MjData(mjm) + mujoco.mj_forward(mjm, mjd) + + # Get reference residuals + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd) + mjw.fwd_position(m, d) + mjw.make_constraint(m, d) + nefc = d.nefc.numpy()[0] + efc_pos_rest = d.efc.pos.numpy()[0, :nefc].copy() + + # Apply uniform translation to all nodes (rigid motion) + mjd2 = mujoco.MjData(mjm) + # All flex nodes have 3 DOFs (slide joints), shift all x by 0.1 + for i in range(0, mjm.nq, 3): + mjd2.qpos[i] += 0.1 # shift x + mujoco.mj_forward(mjm, mjd2) + + d2 = mjw.put_data(mjm, mjd2) + mjw.fwd_position(m, d2) + mjw.make_constraint(m, d2) + nefc2 = d2.nefc.numpy()[0] + efc_pos_shifted = d2.efc.pos.numpy()[0, :nefc2] + + # Residuals should remain near zero for rigid translation + np.testing.assert_allclose( + efc_pos_shifted, efc_pos_rest, atol=1e-4, err_msg="FLEXSTRAIN residuals should be invariant under rigid translation" + ) + + def test_trilinear_gravity_parity(self): + """Test that trilinear simulation matches MuJoCo after multiple steps.""" + mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML) + + # MuJoCo reference + mjd = mujoco.MjData(mjm) + for _ in range(10): + mujoco.mj_step(mjm, mjd) + + # Warp + mjd_warp = mujoco.MjData(mjm) + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd_warp) + z0 = d.qpos.numpy()[0, 2] + for _ in range(10): + mjw.step(m, d) + + qpos_warp = d.qpos.numpy()[0] + qpos_mj = mjd.qpos + + # The cube should have fallen + self.assertLess(qpos_warp[2], z0, "Cube should fall under gravity") + + # Allow larger tolerance for accumulated integration error + np.testing.assert_allclose(qpos_warp, qpos_mj, atol=0.01, err_msg="Trilinear qpos diverges from MuJoCo after 10 steps") + + def test_trilinear_node_positions(self): + """Test that flexnode_xpos are computed correctly from body kinematics.""" + mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML) + mjd = mujoco.MjData(mjm) + mujoco.mj_forward(mjm, mjd) + + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd) + mjw.fwd_position(m, d) + + # Compute expected node positions: xpos_n = body_xpos + body_xmat @ flex_node + nflexnode = mjm.nflexnode + warp_xpos = d.flexnode_xpos.numpy()[0, :nflexnode] + nodeadr = mjm.flex_nodeadr[0] + nodenum = mjm.flex_nodenum[0] + for n in range(nodenum): + bodyid = mjm.flex_nodebodyid[nodeadr + n] + body_xpos = mjd.xpos[bodyid] + body_xmat = mjd.xmat[bodyid].reshape(3, 3) + node_local = mjm.flex_node[nodeadr + n] + expected = body_xpos + body_xmat @ node_local + np.testing.assert_allclose(warp_xpos[n], expected, atol=1e-5, err_msg=f"flexnode_xpos mismatch for node {n}") + + def test_trilinear_passive_forces_parity(self): + """Test passive forces (elasticity) match MuJoCo for trilinear flex.""" + xml = """ + + + + + + + + + """ + mjm = mujoco.MjModel.from_xml_string(xml) + mjd = mujoco.MjData(mjm) + + # perturb first node to generate non-zero elasticity forces + mjd.qpos[0] += 0.01 + mjd.qpos[3] += 0.005 + mujoco.mj_forward(mjm, mjd) + + m = mjw.put_model(mjm) + d = mjw.put_data(mjm, mjd) + + mjw.fwd_position(m, d) + mjw.passive(m, d) + + qfrc_passive_warp = d.qfrc_passive.numpy()[0] + qfrc_passive_mj = mjd.qfrc_passive + + # Verify they match + np.testing.assert_allclose( + qfrc_passive_warp, qfrc_passive_mj, atol=_TOLERANCE, err_msg="qfrc_passive mismatch for trilinear flex with elasticity" + ) + + @parameterized.parameters("strain", "true") + def test_trilinear_equality_types(self, equality): + """Test trilinear with different equality types.""" + xml = f""" + + + + + + + + + + """ + mjm, mjd, m, d = test_data.fixture(xml=xml) + + # Should not crash + mjw.forward(m, d) + + # Constraint count should match + self.assertEqual(d.nefc.numpy()[0], mjd.nefc) + + def test_trilinear_contact_qfrc_constraint(self): + """Test qfrc_constraint parity for trilinear flex with ground contacts.""" + xml = """ + + + """ + mjm, mjd, m, d = test_data.fixture(xml=xml) + mjw.forward(m, d) + + # Verify contacts are generated + nacon = d.nacon.numpy()[0] + self.assertGreater(nacon, 0, "Expected contacts between flex and plane") + self.assertEqual(nacon, mjd.ncon) + + # Verify qfrc_constraint parity + qfrc_warp = d.qfrc_constraint.numpy()[0] + qfrc_mj = mjd.qfrc_constraint + np.testing.assert_allclose(qfrc_warp, qfrc_mj, atol=1e-4, err_msg="qfrc_constraint mismatch for trilinear flex contacts") + + +if __name__ == "__main__": + absltest.main() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py index 7884e999..78c73700 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -77,7 +77,7 @@ class BlockDim: # collision_driver segmented_sort: int = 128 # collision_convex - convex_ccd: int = 256 + convex_ccd: int = 64 # forward actuator_velocity: int = 32 # ray @@ -664,6 +664,7 @@ class EqType(enum.IntEnum): WELD: fix relative position and orientation of two bodies TENDON: couple the lengths of two tendons with cubic FLEX: couple the edge lengths of a flex + FLEXSTRAIN: strain constraint for interpolated flex """ CONNECT = mujoco.mjtEq.mjEQ_CONNECT @@ -671,6 +672,7 @@ class EqType(enum.IntEnum): JOINT = mujoco.mjtEq.mjEQ_JOINT TENDON = mujoco.mjtEq.mjEQ_TENDON FLEX = mujoco.mjtEq.mjEQ_FLEX + FLEXSTRAIN = mujoco.mjtEq.mjEQ_FLEXSTRAIN # unsupported: DISTANCE @@ -974,6 +976,7 @@ class Model: ncam: number of cameras nlight: number of lights nflex: number of flexes + nflexnode: number of nodes in all flexes nflexvert: number of vertices in all flexes nflexedge: number of edges in all flexes nflexelem: number of elements in all flexes @@ -1140,6 +1143,10 @@ class Model: flex_internal: internal collision enabled (nflex,) flex_selfcollide: self-collision mode (nflex,) flex_dim: 1: lines, 2: triangles, 3: tetrahedra (nflex,) + flex_interp: interpolation order (0: vertex, 1+: nodes) (nflex,) + flex_cellnum: cell count per dimension (nflex, 3) + flex_nodeadr: first node address (nflex,) + flex_nodenum: number of nodes (nflex,) flex_vertadr: first vertex address (nflex,) flex_vertnum: number of vertices (nflex,) flex_edgeadr: first edge address (nflex,) @@ -1147,13 +1154,14 @@ class Model: flex_elemadr: first element address (nflex,) flex_elemnum: number of elements (nflex,) flex_elemdataadr: first element vertex id address (nflex,) - flex_stiffnessadr: stiffness matrix address (nflex,) + flex_stiffnessadr: stiffness matrix address (nflex,) flex_elemedgeadr: first element edge id address (nflex,) flex_bendingadr: first bending data address (nflex,) flex_shellnum: number of shells (nflex,) flex_shelldataadr: first shell data address (nflex,) flex_evpairadr: first element-vertex pair address (nflex,) flex_evpairnum: number of element-vertex pairs (nflex,) + flex_nodebodyid: node body ids (nflexnode,) flex_vertbodyid: vertex body ids (nflexvert,) flex_edge: edge vertex ids (2 per edge) (nflexedge, 2) flex_edgeflap: adjacent vertex ids (dim=2 only) (nflexedge, 2) @@ -1162,12 +1170,17 @@ class Model: flex_shell: shell fragment vertex ids (dim per frag) (nflexshelldata,) flex_evpair: element-vertex pair indices (nflexevpair, 2) flex_vert: vertex local positions (nflexvert, 3) + flex_vert0: reference vertex positions in qpos0 (nflexvert, 3) + flex_node: node local positions (nflexnode, 3) + flex_node0: reference node positions in qpos0 (nflexnode, 3) flexedge_length0: edge lengths in qpos0 (nflexedge,) flexedge_invweight0: inv. inertia for the edge (nflexedge,) flex_radius: radius around primitive element (nflex,) flex_stiffness: finite element stiffness matrix (nflexstiffness,) flex_bending: bending stiffness (nflexbending,) flex_damping: Rayleigh's damping coefficient (nflex,) + + flex_edgeequality: edge equality type (0:none,1:edge,2:vert,3:strain) (nflex,) flex_centered: flex vertices are centered at body origin (nflex,) flexedge_J_rownnz: number of nonzeros in Jacobian row (nflexedge,) flexedge_J_rowadr: row start address in colind array (nflexedge,) @@ -1312,6 +1325,7 @@ class Model: nrangefinder: number of rangefinder sensors nmaxcondim: maximum condim across geoms, pairs, and flexes nmaxpyramid: maximum number of pyramid directions + nflexintcell: total interp cells (non-strain) for passive forces nmaxpolygon: maximum number of verts per polygon nmaxmeshdeg: maximum number of polygons per vert is_sparse: constraint Jacobian/Hessian layout (sparse vs dense). Does not affect M, whose @@ -1327,6 +1341,8 @@ class Model: has_fluid: True if wind, density, or viscosity are non-zero at put_model time has_sdf_geom: whether the model contains SDF geoms has_flex_selfcollide: whether any flex has self-collision enabled + has_ellipsoid_geom: whether the model contains ellipsoid geoms + has_3d_flex: whether the model contains 3D flexes max_flex_dim: maximum flex dimension in the model block_dim: block dim options body_tree: list of body ids by tree level @@ -1356,6 +1372,7 @@ class Model: eq_jnt_adr: eq_* addresses of type `JOINT` eq_ten_adr: eq_* addresses of type `TENDON` eq_flex_adr: eq * addresses of type `FLEX + eq_flexstrain_adr: eq_* addresses of type `FLEXSTRAIN` tendon_jnt_adr: joint tendon address tendon_site_pair_adr: site pair tendon address tendon_geom_adr: geom tendon address @@ -1413,13 +1430,18 @@ class Model: M_mulm_col: sparse matmul column indices M_mulm_madr: sparse matmul matrix addresses flexelem_geom_pair_filtered: conaffinity-filtered element vs geom pairs (*, 2) - flexshell_geom_pair_filtered: conaffinity-filtered shell vs geom pairs (*, 2) flexvert_geom_pair_filtered: conaffinity-filtered vertex vs geom pairs (*, 2) flex_elemflexid: maps each element index directly to its flexid (nflexelem,) flex_shellflexid: maps each shell index directly to its flexid (nflexshelldata,) flex_evpairflexid: maps each element-vertex pair directly to its flexid (nflexevpair,) flex_vertflexid: maps each vertex index directly to its flexid (nflexvert,) flex_shelladr: maps each flex to its start shell index (nflex,) + flex_cell_map: precomputed flex cell mapping (nflexintcell,) + flexstrain_J_rownnz: number of nonzeros in flex strain Jacobian row (neq_flexstrain,) + flexstrain_J_rowadr: row start address in colind array (neq_flexstrain,) + flexstrain_J_colind: column indices in sparse flex strain Jacobian (nJfs,) + neq_flexstrain: number of flex strain equality constraints + nJfs: number of non-zeros in sparse flex strain Jacobian """ nq: int @@ -1438,6 +1460,7 @@ class Model: ncam: int nlight: int nflex: int + nflexnode: int nflexvert: int nflexedge: int nflexelem: int @@ -1604,6 +1627,10 @@ class Model: flex_internal: array("nflex", int) flex_selfcollide: array("nflex", int) flex_dim: array("nflex", int) + flex_interp: array("nflex", int) + flex_cellnum: array("nflex", wp.vec3i) + flex_nodeadr: array("nflex", int) + flex_nodenum: array("nflex", int) flex_vertadr: array("nflex", int) flex_vertnum: array("nflex", int) flex_edgeadr: array("nflex", int) @@ -1618,6 +1645,7 @@ class Model: flex_shelldataadr: array("nflex", int) flex_evpairadr: array("nflex", int) flex_evpairnum: array("nflex", int) + flex_nodebodyid: array("nflexnode", int) flex_vertbodyid: array("nflexvert", int) flex_edge: array("nflexedge", wp.vec2i) flex_edgeflap: array("nflexedge", wp.vec2i) @@ -1626,12 +1654,16 @@ class Model: flex_shell: array("nflexshelldata", int) flex_evpair: array("nflexevpair", wp.vec2i) flex_vert: array("nflexvert", wp.vec3) + flex_vert0: array("nflexvert", wp.vec3) + flex_node: array("nflexnode", wp.vec3) + flex_node0: array("nflexnode", wp.vec3) flexedge_length0: array("nflexedge", float) flexedge_invweight0: array("nflexedge", float) flex_radius: array("nflex", float) flex_stiffness: array("nflexstiffness", float) flex_bending: array("nflexbending", float) flex_damping: array("nflex", float) + flex_edgeequality: array("nflex", int) flex_centered: array("nflex", bool) flexedge_J_rownnz: array("nflexedge", int) flexedge_J_rowadr: array("nflexedge", int) @@ -1774,6 +1806,7 @@ class Model: nrangefinder: int nmaxcondim: int nmaxpyramid: int + nflexintcell: int nmaxpolygon: int nmaxmeshdeg: int is_sparse: bool @@ -1788,6 +1821,8 @@ class Model: has_fluid: bool has_sdf_geom: bool has_flex_selfcollide: bool + has_ellipsoid_geom: bool + has_3d_flex: bool max_flex_dim: int block_dim: BlockDim body_tree: tuple[wp.array[int], ...] @@ -1813,6 +1848,7 @@ class Model: eq_jnt_adr: wp.array[int] eq_ten_adr: wp.array[int] eq_flex_adr: wp.array[int] + eq_flexstrain_adr: wp.array[int] tendon_jnt_adr: wp.array[int] tendon_site_pair_adr: wp.array[int] tendon_geom_adr: wp.array[int] @@ -1868,13 +1904,18 @@ class Model: M_mulm_col: wp.array[int] # column index to gather from M_mulm_madr: wp.array[int] # matrix address to read flexelem_geom_pair_filtered: wp.array[wp.vec2i] - flexshell_geom_pair_filtered: wp.array[wp.vec2i] flexvert_geom_pair_filtered: wp.array[wp.vec2i] flex_elemflexid: array("nflexelem", int) flex_shellflexid: array("nflexshelldata", int) flex_evpairflexid: array("nflexevpair", int) flex_vertflexid: array("nflexvert", int) flex_shelladr: array("nflex", int) + flex_cell_map: array("nflexintcell", wp.vec4i) + flexstrain_J_rownnz: array("neq_flexstrain", int) + flexstrain_J_rowadr: array("neq_flexstrain", int) + flexstrain_J_colind: array("nJfs", int) + neq_flexstrain: int + nJfs: int class ContactType(enum.IntFlag): @@ -2136,6 +2177,7 @@ class Data: ncollision: collision count from broadphase (1,) flex_aabb_min: dynamic flex object bounding box min (nworld, nflex, 3) flex_aabb_max: dynamic flex object bounding box max (nworld, nflex, 3) + flexnode_xpos: cartesian flex node positions (nworld, nflexnode, 3) overflow: overflow bitmask (OverflowType) (nworld,) """ @@ -2279,6 +2321,7 @@ class Data: ncollision: array(1, int) flex_aabb_min: array("nworld", "nflex", wp.vec3) flex_aabb_max: array("nworld", "nflex", wp.vec3) + flexnode_xpos: array("nworld", "nflexnode", wp.vec3) overflow: array("nworld", int) @@ -2308,6 +2351,7 @@ class SolverContext: jv: wp.array2d[float] quad: wp.array2d[wp.vec3] alpha: wp.array[float] + grad_scale: wp.array[float] improvement: wp.array[float] prev_grad: wp.array2d[float] prev_Mgrad: wp.array2d[float] diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 4c061f64..e5c95947 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -72,10 +72,6 @@ def _collision_shim( flex_priority: wp.array[int], flex_radius: wp.array[float], flex_selfcollide: wp.array[int], - flex_shell: wp.array[int], - flex_shelladr: wp.array[int], - flex_shelldataadr: wp.array[int], - flex_shellflexid: wp.array[int], flex_solimp: wp.array[mjwp_types.vec5], flex_solmix: wp.array[float], flex_solref: wp.array[wp.vec2], @@ -84,7 +80,6 @@ def _collision_shim( flex_vertflexid: wp.array[int], flex_vertnum: wp.array[int], flexelem_geom_pair_filtered: wp.array[wp.vec2i], - flexshell_geom_pair_filtered: wp.array[wp.vec2i], flexvert_geom_pair_filtered: wp.array[wp.vec2i], geom_aabb: wp.array3d[wp.vec3], geom_bodyid: wp.array[int], @@ -104,6 +99,8 @@ def _collision_shim( geom_solmix: wp.array2d[float], geom_solref: wp.array2d[wp.vec2], geom_type: wp.array[int], + has_3d_flex: bool, + has_ellipsoid_geom: bool, has_flex_selfcollide: bool, has_sdf_geom: bool, hfield_adr: wp.array[int], @@ -221,10 +218,6 @@ def _collision_shim( _m.flex_priority = flex_priority _m.flex_radius = flex_radius _m.flex_selfcollide = flex_selfcollide - _m.flex_shell = flex_shell - _m.flex_shelladr = flex_shelladr - _m.flex_shelldataadr = flex_shelldataadr - _m.flex_shellflexid = flex_shellflexid _m.flex_solimp = flex_solimp _m.flex_solmix = flex_solmix _m.flex_solref = flex_solref @@ -233,7 +226,6 @@ def _collision_shim( _m.flex_vertflexid = flex_vertflexid _m.flex_vertnum = flex_vertnum _m.flexelem_geom_pair_filtered = flexelem_geom_pair_filtered - _m.flexshell_geom_pair_filtered = flexshell_geom_pair_filtered _m.flexvert_geom_pair_filtered = flexvert_geom_pair_filtered _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid @@ -253,6 +245,8 @@ def _collision_shim( _m.geom_solmix = geom_solmix _m.geom_solref = geom_solref _m.geom_type = geom_type + _m.has_3d_flex = has_3d_flex + _m.has_ellipsoid_geom = has_ellipsoid_geom _m.has_flex_selfcollide = has_flex_selfcollide _m.has_sdf_geom = has_sdf_geom _m.hfield_adr = hfield_adr @@ -447,10 +441,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data): m._impl.flex_priority, m._impl.flex_radius, m._impl.flex_selfcollide, - m._impl.flex_shell, - m._impl.flex_shelladr, - m._impl.flex_shelldataadr, - m._impl.flex_shellflexid, m._impl.flex_solimp, m._impl.flex_solmix, m._impl.flex_solref, @@ -459,7 +449,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data): m._impl.flex_vertflexid, m.flex_vertnum, m._impl.flexelem_geom_pair_filtered, - m._impl.flexshell_geom_pair_filtered, m._impl.flexvert_geom_pair_filtered, m.geom_aabb, m.geom_bodyid, @@ -479,6 +468,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data): m.geom_solmix, m.geom_solref, m.geom_type, + m._impl.has_3d_flex, + m._impl.has_ellipsoid_geom, m._impl.has_flex_selfcollide, m._impl.has_sdf_geom, m.hfield_adr, diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 9440af51..a68351ff 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -140,6 +140,7 @@ def _forward_shim( eq_connect_adr: wp.array[int], eq_data: wp.array2d[mjwp_types.vec11], eq_flex_adr: wp.array[int], + eq_flexstrain_adr: wp.array[int], eq_jnt_adr: wp.array[int], eq_obj1id: wp.array[int], eq_obj2id: wp.array[int], @@ -151,6 +152,8 @@ def _forward_shim( eq_wld_adr: wp.array[int], flex_bending: wp.array[float], flex_bendingadr: wp.array[int], + flex_cell_map: wp.array[wp.vec4i], + flex_cellnum: wp.array[wp.vec3i], flex_centered: wp.array[bool], flex_conaffinity: wp.array[int], flex_condim: wp.array[int], @@ -159,6 +162,7 @@ def _forward_shim( flex_dim: wp.array[int], flex_edge: wp.array[wp.vec2i], flex_edgeadr: wp.array[int], + flex_edgeequality: wp.array[int], flex_edgeflap: wp.array[wp.vec2i], flex_edgenum: wp.array[int], flex_elem: wp.array[int], @@ -175,20 +179,25 @@ def _forward_shim( flex_friction: wp.array[wp.vec3], flex_gap: wp.array[float], flex_internal: wp.array[int], + flex_interp: wp.array[int], flex_margin: wp.array[float], + flex_node: wp.array[wp.vec3], + flex_node0: wp.array[wp.vec3], + flex_nodeadr: wp.array[int], + flex_nodebodyid: wp.array[int], + flex_nodenum: wp.array[int], flex_priority: wp.array[int], flex_radius: wp.array[float], flex_selfcollide: wp.array[int], flex_shell: wp.array[int], - flex_shelladr: wp.array[int], flex_shelldataadr: wp.array[int], - flex_shellflexid: wp.array[int], flex_solimp: wp.array[mjwp_types.vec5], flex_solmix: wp.array[float], flex_solref: wp.array[wp.vec2], flex_stiffness: wp.array[float], flex_stiffnessadr: wp.array[int], flex_vert: wp.array[wp.vec3], + flex_vert0: wp.array[wp.vec3], flex_vertadr: wp.array[int], flex_vertbodyid: wp.array[int], flex_vertflexid: wp.array[int], @@ -199,7 +208,9 @@ def _forward_shim( flexedge_invweight0: wp.array[float], flexedge_length0: wp.array[float], flexelem_geom_pair_filtered: wp.array[wp.vec2i], - flexshell_geom_pair_filtered: wp.array[wp.vec2i], + flexstrain_J_colind: wp.array[int], + flexstrain_J_rowadr: wp.array[int], + flexstrain_J_rownnz: wp.array[int], flexvert_geom_pair_filtered: wp.array[wp.vec2i], geom_aabb: wp.array3d[wp.vec3], geom_bodyid: wp.array[int], @@ -225,6 +236,8 @@ def _forward_shim( geom_solmix: wp.array2d[float], geom_solref: wp.array2d[wp.vec2], geom_type: wp.array[int], + has_3d_flex: bool, + has_ellipsoid_geom: bool, has_flex_selfcollide: bool, has_fluid: bool, has_sdf_geom: bool, @@ -296,6 +309,8 @@ def _forward_shim( nflexedge: int, nflexelem: int, nflexevpair: int, + nflexintcell: int, + nflexnode: int, nflexvert: int, ngeom: int, nhistory: int, @@ -502,6 +517,7 @@ def _forward_shim( flexedge_J: wp.array2d[float], flexedge_length: wp.array2d[float], flexedge_velocity: wp.array2d[float], + flexnode_xpos: wp.array2d[wp.vec3], flexvert_xpos: wp.array2d[wp.vec3], geom_xmat: wp.array2d[wp.mat33], geom_xpos: wp.array2d[wp.vec3], @@ -712,6 +728,7 @@ def _forward_shim( _m.eq_connect_adr = eq_connect_adr _m.eq_data = eq_data _m.eq_flex_adr = eq_flex_adr + _m.eq_flexstrain_adr = eq_flexstrain_adr _m.eq_jnt_adr = eq_jnt_adr _m.eq_obj1id = eq_obj1id _m.eq_obj2id = eq_obj2id @@ -723,6 +740,8 @@ def _forward_shim( _m.eq_wld_adr = eq_wld_adr _m.flex_bending = flex_bending _m.flex_bendingadr = flex_bendingadr + _m.flex_cell_map = flex_cell_map + _m.flex_cellnum = flex_cellnum _m.flex_centered = flex_centered _m.flex_conaffinity = flex_conaffinity _m.flex_condim = flex_condim @@ -731,6 +750,7 @@ def _forward_shim( _m.flex_dim = flex_dim _m.flex_edge = flex_edge _m.flex_edgeadr = flex_edgeadr + _m.flex_edgeequality = flex_edgeequality _m.flex_edgeflap = flex_edgeflap _m.flex_edgenum = flex_edgenum _m.flex_elem = flex_elem @@ -747,20 +767,25 @@ def _forward_shim( _m.flex_friction = flex_friction _m.flex_gap = flex_gap _m.flex_internal = flex_internal + _m.flex_interp = flex_interp _m.flex_margin = flex_margin + _m.flex_node = flex_node + _m.flex_node0 = flex_node0 + _m.flex_nodeadr = flex_nodeadr + _m.flex_nodebodyid = flex_nodebodyid + _m.flex_nodenum = flex_nodenum _m.flex_priority = flex_priority _m.flex_radius = flex_radius _m.flex_selfcollide = flex_selfcollide _m.flex_shell = flex_shell - _m.flex_shelladr = flex_shelladr _m.flex_shelldataadr = flex_shelldataadr - _m.flex_shellflexid = flex_shellflexid _m.flex_solimp = flex_solimp _m.flex_solmix = flex_solmix _m.flex_solref = flex_solref _m.flex_stiffness = flex_stiffness _m.flex_stiffnessadr = flex_stiffnessadr _m.flex_vert = flex_vert + _m.flex_vert0 = flex_vert0 _m.flex_vertadr = flex_vertadr _m.flex_vertbodyid = flex_vertbodyid _m.flex_vertflexid = flex_vertflexid @@ -771,7 +796,9 @@ def _forward_shim( _m.flexedge_invweight0 = flexedge_invweight0 _m.flexedge_length0 = flexedge_length0 _m.flexelem_geom_pair_filtered = flexelem_geom_pair_filtered - _m.flexshell_geom_pair_filtered = flexshell_geom_pair_filtered + _m.flexstrain_J_colind = flexstrain_J_colind + _m.flexstrain_J_rowadr = flexstrain_J_rowadr + _m.flexstrain_J_rownnz = flexstrain_J_rownnz _m.flexvert_geom_pair_filtered = flexvert_geom_pair_filtered _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid @@ -797,6 +824,8 @@ def _forward_shim( _m.geom_solmix = geom_solmix _m.geom_solref = geom_solref _m.geom_type = geom_type + _m.has_3d_flex = has_3d_flex + _m.has_ellipsoid_geom = has_ellipsoid_geom _m.has_flex_selfcollide = has_flex_selfcollide _m.has_fluid = has_fluid _m.has_sdf_geom = has_sdf_geom @@ -868,6 +897,8 @@ def _forward_shim( _m.nflexedge = nflexedge _m.nflexelem = nflexelem _m.nflexevpair = nflexevpair + _m.nflexintcell = nflexintcell + _m.nflexnode = nflexnode _m.nflexvert = nflexvert _m.ngeom = ngeom _m.nhistory = nhistory @@ -1104,6 +1135,7 @@ def _forward_shim( _d.flexedge_J = flexedge_J _d.flexedge_length = flexedge_length _d.flexedge_velocity = flexedge_velocity + _d.flexnode_xpos = flexnode_xpos _d.flexvert_xpos = flexvert_xpos _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos @@ -1231,6 +1263,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, + 'flexnode_xpos': d._impl.flexnode_xpos.shape, 'flexvert_xpos': d._impl.flexvert_xpos.shape, 'geom_xmat': d.geom_xmat.shape, 'geom_xpos': d.geom_xpos.shape, @@ -1342,7 +1375,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _forward_shim, - num_outputs=145, + num_outputs=146, output_dims=output_dims, vmap_method=None, in_out_argnames=set([ @@ -1383,6 +1416,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'flexedge_J', 'flexedge_length', 'flexedge_velocity', + 'flexnode_xpos', 'flexvert_xpos', 'geom_xmat', 'geom_xpos', @@ -1765,6 +1799,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.eq_connect_adr, m.eq_data, m._impl.eq_flex_adr, + m._impl.eq_flexstrain_adr, m._impl.eq_jnt_adr, m.eq_obj1id, m.eq_obj2id, @@ -1776,6 +1811,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.eq_wld_adr, m._impl.flex_bending, m._impl.flex_bendingadr, + m._impl.flex_cell_map, + m._impl.flex_cellnum, m._impl.flex_centered, m._impl.flex_conaffinity, m._impl.flex_condim, @@ -1784,6 +1821,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.flex_dim, m._impl.flex_edge, m._impl.flex_edgeadr, + m._impl.flex_edgeequality, m._impl.flex_edgeflap, m._impl.flex_edgenum, m._impl.flex_elem, @@ -1800,20 +1838,25 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.flex_friction, m._impl.flex_gap, m._impl.flex_internal, + m.flex_interp, m._impl.flex_margin, + m._impl.flex_node, + m.flex_node0, + m.flex_nodeadr, + m.flex_nodebodyid, + m.flex_nodenum, m._impl.flex_priority, m._impl.flex_radius, m._impl.flex_selfcollide, m._impl.flex_shell, - m._impl.flex_shelladr, m._impl.flex_shelldataadr, - m._impl.flex_shellflexid, m._impl.flex_solimp, m._impl.flex_solmix, m._impl.flex_solref, m._impl.flex_stiffness, m._impl.flex_stiffnessadr, m._impl.flex_vert, + m.flex_vert0, m.flex_vertadr, m._impl.flex_vertbodyid, m._impl.flex_vertflexid, @@ -1824,7 +1867,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.flexedge_invweight0, m._impl.flexedge_length0, m._impl.flexelem_geom_pair_filtered, - m._impl.flexshell_geom_pair_filtered, + m._impl.flexstrain_J_colind, + m._impl.flexstrain_J_rowadr, + m._impl.flexstrain_J_rownnz, m._impl.flexvert_geom_pair_filtered, m.geom_aabb, m.geom_bodyid, @@ -1850,6 +1895,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.geom_solmix, m.geom_solref, m.geom_type, + m._impl.has_3d_flex, + m._impl.has_ellipsoid_geom, m._impl.has_flex_selfcollide, m._impl.has_fluid, m._impl.has_sdf_geom, @@ -1921,6 +1968,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.nflexedge, m._impl.nflexelem, m._impl.nflexevpair, + m._impl.nflexintcell, + m._impl.nflexnode, m._impl.nflexvert, m.ngeom, m.nhistory, @@ -2126,6 +2175,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.flexedge_J, d._impl.flexedge_length, d._impl.flexedge_velocity, + d._impl.flexnode_xpos, d._impl.flexvert_xpos, d.geom_xmat, d.geom_xpos, @@ -2280,114 +2330,115 @@ def _forward_jax_impl(m: types.Model, d: types.Data): '_impl.flexedge_J': out[34], '_impl.flexedge_length': out[35], '_impl.flexedge_velocity': out[36], - '_impl.flexvert_xpos': out[37], - 'geom_xmat': out[38], - 'geom_xpos': out[39], - 'history': out[40], - '_impl.island_dofadr': out[41], - '_impl.island_idofadr': out[42], - '_impl.island_iefcadr': out[43], - '_impl.island_ne': out[44], - '_impl.island_nefc': out[45], - '_impl.island_nf': out[46], - '_impl.island_nv': out[47], - '_impl.light_xdir': out[48], - '_impl.light_xpos': out[49], - '_impl.map_dof2idof': out[50], - '_impl.map_efc2iefc': out[51], - '_impl.map_idof2dof': out[52], - '_impl.map_iefc2efc': out[53], - '_impl.moment_colind': out[54], - '_impl.moment_rowadr': out[55], - '_impl.moment_rownnz': out[56], - '_impl.nacon': out[57], - '_impl.nbody_awake': out[58], - '_impl.ncdof': out[59], - '_impl.ncollision': out[60], - '_impl.ne': out[61], - '_impl.nefc': out[62], - '_impl.nf': out[63], - '_impl.nidof': out[64], - '_impl.nisland': out[65], - '_impl.nl': out[66], - '_impl.ntree_awake': out[67], - '_impl.nv_awake': out[68], - '_impl.overflow': out[69], - '_impl.qLD': out[70], - '_impl.qLDiagInv': out[71], - 'qacc': out[72], - 'qacc_smooth': out[73], - 'qfrc_actuator': out[74], - 'qfrc_bias': out[75], - 'qfrc_constraint': out[76], - '_impl.qfrc_damper': out[77], - 'qfrc_fluid': out[78], - 'qfrc_gravcomp': out[79], - 'qfrc_passive': out[80], - 'qfrc_smooth': out[81], - '_impl.qfrc_spring': out[82], - 'qvel': out[83], - 'sensordata': out[84], - 'site_xmat': out[85], - 'site_xpos': out[86], - '_impl.solver_niter': out[87], - '_impl.subtree_angmom': out[88], - 'subtree_com': out[89], - '_impl.subtree_linvel': out[90], - '_impl.ten_J': out[91], - 'ten_length': out[92], - '_impl.ten_velocity': out[93], - '_impl.ten_wrapadr': out[94], - '_impl.ten_wrapnum': out[95], - '_impl.tree_asleep': out[96], - '_impl.tree_awake': out[97], - '_impl.tree_island': out[98], - '_impl.wrap_obj': out[99], - '_impl.wrap_xpos': out[100], - 'xanchor': out[101], - 'xaxis': out[102], - 'ximat': out[103], - 'xipos': out[104], - 'xmat': out[105], - 'xpos': out[106], - 'xquat': out[107], - '_impl.contact__dim': out[108], - '_impl.contact__dist': out[109], - '_impl.contact__efc_address': out[110], - '_impl.contact__elem': out[111], - '_impl.contact__flex': out[112], - '_impl.contact__frame': out[113], - '_impl.contact__friction': out[114], - '_impl.contact__geom': out[115], - '_impl.contact__geomcollisionid': out[116], - '_impl.contact__includemargin': out[117], - '_impl.contact__pos': out[118], - '_impl.contact__solimp': out[119], - '_impl.contact__solref': out[120], - '_impl.contact__solreffriction': out[121], - '_impl.contact__type': out[122], - '_impl.contact__vert': out[123], - '_impl.contact__worldid': out[124], - '_impl.efc__D': out[125], - '_impl.efc__J': out[126], - '_impl.efc__J_colind': out[127], - '_impl.efc__J_rowadr': out[128], - '_impl.efc__J_rownnz': out[129], - '_impl.efc__Jqvel': out[130], - '_impl.efc__Ma': out[131], - '_impl.efc__aref': out[132], - '_impl.efc__force': out[133], - '_impl.efc__frictionloss': out[134], - '_impl.efc__id': out[135], - '_impl.efc__island': out[136], - '_impl.efc__jtdaj_adr': out[137], - '_impl.efc__jtdaj_nblock': out[138], - '_impl.efc__jtdaj_nrow': out[139], - '_impl.efc__margin': out[140], - '_impl.efc__pos': out[141], - '_impl.efc__state': out[142], - '_impl.efc__type': out[143], - '_impl.efc__vel': out[144], + '_impl.flexnode_xpos': out[37], + '_impl.flexvert_xpos': out[38], + 'geom_xmat': out[39], + 'geom_xpos': out[40], + 'history': out[41], + '_impl.island_dofadr': out[42], + '_impl.island_idofadr': out[43], + '_impl.island_iefcadr': out[44], + '_impl.island_ne': out[45], + '_impl.island_nefc': out[46], + '_impl.island_nf': out[47], + '_impl.island_nv': out[48], + '_impl.light_xdir': out[49], + '_impl.light_xpos': out[50], + '_impl.map_dof2idof': out[51], + '_impl.map_efc2iefc': out[52], + '_impl.map_idof2dof': out[53], + '_impl.map_iefc2efc': out[54], + '_impl.moment_colind': out[55], + '_impl.moment_rowadr': out[56], + '_impl.moment_rownnz': out[57], + '_impl.nacon': out[58], + '_impl.nbody_awake': out[59], + '_impl.ncdof': out[60], + '_impl.ncollision': out[61], + '_impl.ne': out[62], + '_impl.nefc': out[63], + '_impl.nf': out[64], + '_impl.nidof': out[65], + '_impl.nisland': out[66], + '_impl.nl': out[67], + '_impl.ntree_awake': out[68], + '_impl.nv_awake': out[69], + '_impl.overflow': out[70], + '_impl.qLD': out[71], + '_impl.qLDiagInv': out[72], + 'qacc': out[73], + 'qacc_smooth': out[74], + 'qfrc_actuator': out[75], + 'qfrc_bias': out[76], + 'qfrc_constraint': out[77], + '_impl.qfrc_damper': out[78], + 'qfrc_fluid': out[79], + 'qfrc_gravcomp': out[80], + 'qfrc_passive': out[81], + 'qfrc_smooth': out[82], + '_impl.qfrc_spring': out[83], + 'qvel': out[84], + 'sensordata': out[85], + 'site_xmat': out[86], + 'site_xpos': out[87], + '_impl.solver_niter': out[88], + '_impl.subtree_angmom': out[89], + 'subtree_com': out[90], + '_impl.subtree_linvel': out[91], + '_impl.ten_J': out[92], + 'ten_length': out[93], + '_impl.ten_velocity': out[94], + '_impl.ten_wrapadr': out[95], + '_impl.ten_wrapnum': out[96], + '_impl.tree_asleep': out[97], + '_impl.tree_awake': out[98], + '_impl.tree_island': out[99], + '_impl.wrap_obj': out[100], + '_impl.wrap_xpos': out[101], + 'xanchor': out[102], + 'xaxis': out[103], + 'ximat': out[104], + 'xipos': out[105], + 'xmat': out[106], + 'xpos': out[107], + 'xquat': out[108], + '_impl.contact__dim': out[109], + '_impl.contact__dist': out[110], + '_impl.contact__efc_address': out[111], + '_impl.contact__elem': out[112], + '_impl.contact__flex': out[113], + '_impl.contact__frame': out[114], + '_impl.contact__friction': out[115], + '_impl.contact__geom': out[116], + '_impl.contact__geomcollisionid': out[117], + '_impl.contact__includemargin': out[118], + '_impl.contact__pos': out[119], + '_impl.contact__solimp': out[120], + '_impl.contact__solref': out[121], + '_impl.contact__solreffriction': out[122], + '_impl.contact__type': out[123], + '_impl.contact__vert': out[124], + '_impl.contact__worldid': out[125], + '_impl.efc__D': out[126], + '_impl.efc__J': out[127], + '_impl.efc__J_colind': out[128], + '_impl.efc__J_rowadr': out[129], + '_impl.efc__J_rownnz': out[130], + '_impl.efc__Jqvel': out[131], + '_impl.efc__Ma': out[132], + '_impl.efc__aref': out[133], + '_impl.efc__force': out[134], + '_impl.efc__frictionloss': out[135], + '_impl.efc__id': out[136], + '_impl.efc__island': out[137], + '_impl.efc__jtdaj_adr': out[138], + '_impl.efc__jtdaj_nblock': out[139], + '_impl.efc__jtdaj_nrow': out[140], + '_impl.efc__margin': out[141], + '_impl.efc__pos': out[142], + '_impl.efc__state': out[143], + '_impl.efc__type': out[144], + '_impl.efc__vel': out[145], }) return d @@ -2506,6 +2557,7 @@ def _step_shim( eq_connect_adr: wp.array[int], eq_data: wp.array2d[mjwp_types.vec11], eq_flex_adr: wp.array[int], + eq_flexstrain_adr: wp.array[int], eq_jnt_adr: wp.array[int], eq_obj1id: wp.array[int], eq_obj2id: wp.array[int], @@ -2517,6 +2569,8 @@ def _step_shim( eq_wld_adr: wp.array[int], flex_bending: wp.array[float], flex_bendingadr: wp.array[int], + flex_cell_map: wp.array[wp.vec4i], + flex_cellnum: wp.array[wp.vec3i], flex_centered: wp.array[bool], flex_conaffinity: wp.array[int], flex_condim: wp.array[int], @@ -2525,6 +2579,7 @@ def _step_shim( flex_dim: wp.array[int], flex_edge: wp.array[wp.vec2i], flex_edgeadr: wp.array[int], + flex_edgeequality: wp.array[int], flex_edgeflap: wp.array[wp.vec2i], flex_edgenum: wp.array[int], flex_elem: wp.array[int], @@ -2541,20 +2596,25 @@ def _step_shim( flex_friction: wp.array[wp.vec3], flex_gap: wp.array[float], flex_internal: wp.array[int], + flex_interp: wp.array[int], flex_margin: wp.array[float], + flex_node: wp.array[wp.vec3], + flex_node0: wp.array[wp.vec3], + flex_nodeadr: wp.array[int], + flex_nodebodyid: wp.array[int], + flex_nodenum: wp.array[int], flex_priority: wp.array[int], flex_radius: wp.array[float], flex_selfcollide: wp.array[int], flex_shell: wp.array[int], - flex_shelladr: wp.array[int], flex_shelldataadr: wp.array[int], - flex_shellflexid: wp.array[int], flex_solimp: wp.array[mjwp_types.vec5], flex_solmix: wp.array[float], flex_solref: wp.array[wp.vec2], flex_stiffness: wp.array[float], flex_stiffnessadr: wp.array[int], flex_vert: wp.array[wp.vec3], + flex_vert0: wp.array[wp.vec3], flex_vertadr: wp.array[int], flex_vertbodyid: wp.array[int], flex_vertflexid: wp.array[int], @@ -2565,7 +2625,9 @@ def _step_shim( flexedge_invweight0: wp.array[float], flexedge_length0: wp.array[float], flexelem_geom_pair_filtered: wp.array[wp.vec2i], - flexshell_geom_pair_filtered: wp.array[wp.vec2i], + flexstrain_J_colind: wp.array[int], + flexstrain_J_rowadr: wp.array[int], + flexstrain_J_rownnz: wp.array[int], flexvert_geom_pair_filtered: wp.array[wp.vec2i], geom_aabb: wp.array3d[wp.vec3], geom_bodyid: wp.array[int], @@ -2591,6 +2653,8 @@ def _step_shim( geom_solmix: wp.array2d[float], geom_solref: wp.array2d[wp.vec2], geom_type: wp.array[int], + has_3d_flex: bool, + has_ellipsoid_geom: bool, has_flex_selfcollide: bool, has_fluid: bool, has_sdf_geom: bool, @@ -2665,6 +2729,8 @@ def _step_shim( nflexedge: int, nflexelem: int, nflexevpair: int, + nflexintcell: int, + nflexnode: int, nflexvert: int, ngeom: int, nhistory: int, @@ -2875,6 +2941,7 @@ def _step_shim( flexedge_J: wp.array2d[float], flexedge_length: wp.array2d[float], flexedge_velocity: wp.array2d[float], + flexnode_xpos: wp.array2d[wp.vec3], flexvert_xpos: wp.array2d[wp.vec3], geom_xmat: wp.array2d[wp.mat33], geom_xpos: wp.array2d[wp.vec3], @@ -3094,6 +3161,7 @@ def _step_shim( _m.eq_connect_adr = eq_connect_adr _m.eq_data = eq_data _m.eq_flex_adr = eq_flex_adr + _m.eq_flexstrain_adr = eq_flexstrain_adr _m.eq_jnt_adr = eq_jnt_adr _m.eq_obj1id = eq_obj1id _m.eq_obj2id = eq_obj2id @@ -3105,6 +3173,8 @@ def _step_shim( _m.eq_wld_adr = eq_wld_adr _m.flex_bending = flex_bending _m.flex_bendingadr = flex_bendingadr + _m.flex_cell_map = flex_cell_map + _m.flex_cellnum = flex_cellnum _m.flex_centered = flex_centered _m.flex_conaffinity = flex_conaffinity _m.flex_condim = flex_condim @@ -3113,6 +3183,7 @@ def _step_shim( _m.flex_dim = flex_dim _m.flex_edge = flex_edge _m.flex_edgeadr = flex_edgeadr + _m.flex_edgeequality = flex_edgeequality _m.flex_edgeflap = flex_edgeflap _m.flex_edgenum = flex_edgenum _m.flex_elem = flex_elem @@ -3129,20 +3200,25 @@ def _step_shim( _m.flex_friction = flex_friction _m.flex_gap = flex_gap _m.flex_internal = flex_internal + _m.flex_interp = flex_interp _m.flex_margin = flex_margin + _m.flex_node = flex_node + _m.flex_node0 = flex_node0 + _m.flex_nodeadr = flex_nodeadr + _m.flex_nodebodyid = flex_nodebodyid + _m.flex_nodenum = flex_nodenum _m.flex_priority = flex_priority _m.flex_radius = flex_radius _m.flex_selfcollide = flex_selfcollide _m.flex_shell = flex_shell - _m.flex_shelladr = flex_shelladr _m.flex_shelldataadr = flex_shelldataadr - _m.flex_shellflexid = flex_shellflexid _m.flex_solimp = flex_solimp _m.flex_solmix = flex_solmix _m.flex_solref = flex_solref _m.flex_stiffness = flex_stiffness _m.flex_stiffnessadr = flex_stiffnessadr _m.flex_vert = flex_vert + _m.flex_vert0 = flex_vert0 _m.flex_vertadr = flex_vertadr _m.flex_vertbodyid = flex_vertbodyid _m.flex_vertflexid = flex_vertflexid @@ -3153,7 +3229,9 @@ def _step_shim( _m.flexedge_invweight0 = flexedge_invweight0 _m.flexedge_length0 = flexedge_length0 _m.flexelem_geom_pair_filtered = flexelem_geom_pair_filtered - _m.flexshell_geom_pair_filtered = flexshell_geom_pair_filtered + _m.flexstrain_J_colind = flexstrain_J_colind + _m.flexstrain_J_rowadr = flexstrain_J_rowadr + _m.flexstrain_J_rownnz = flexstrain_J_rownnz _m.flexvert_geom_pair_filtered = flexvert_geom_pair_filtered _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid @@ -3179,6 +3257,8 @@ def _step_shim( _m.geom_solmix = geom_solmix _m.geom_solref = geom_solref _m.geom_type = geom_type + _m.has_3d_flex = has_3d_flex + _m.has_ellipsoid_geom = has_ellipsoid_geom _m.has_flex_selfcollide = has_flex_selfcollide _m.has_fluid = has_fluid _m.has_sdf_geom = has_sdf_geom @@ -3253,6 +3333,8 @@ def _step_shim( _m.nflexedge = nflexedge _m.nflexelem = nflexelem _m.nflexevpair = nflexevpair + _m.nflexintcell = nflexintcell + _m.nflexnode = nflexnode _m.nflexvert = nflexvert _m.ngeom = ngeom _m.nhistory = nhistory @@ -3493,6 +3575,7 @@ def _step_shim( _d.flexedge_J = flexedge_J _d.flexedge_length = flexedge_length _d.flexedge_velocity = flexedge_velocity + _d.flexnode_xpos = flexnode_xpos _d.flexvert_xpos = flexvert_xpos _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos @@ -3622,6 +3705,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, + 'flexnode_xpos': d._impl.flexnode_xpos.shape, 'flexvert_xpos': d._impl.flexvert_xpos.shape, 'geom_xmat': d.geom_xmat.shape, 'geom_xpos': d.geom_xpos.shape, @@ -3737,7 +3821,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _step_shim, - num_outputs=150, + num_outputs=151, output_dims=output_dims, vmap_method=None, in_out_argnames=set([ @@ -3779,6 +3863,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'flexedge_J', 'flexedge_length', 'flexedge_velocity', + 'flexnode_xpos', 'flexvert_xpos', 'geom_xmat', 'geom_xpos', @@ -4177,6 +4262,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.eq_connect_adr, m.eq_data, m._impl.eq_flex_adr, + m._impl.eq_flexstrain_adr, m._impl.eq_jnt_adr, m.eq_obj1id, m.eq_obj2id, @@ -4188,6 +4274,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.eq_wld_adr, m._impl.flex_bending, m._impl.flex_bendingadr, + m._impl.flex_cell_map, + m._impl.flex_cellnum, m._impl.flex_centered, m._impl.flex_conaffinity, m._impl.flex_condim, @@ -4196,6 +4284,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.flex_dim, m._impl.flex_edge, m._impl.flex_edgeadr, + m._impl.flex_edgeequality, m._impl.flex_edgeflap, m._impl.flex_edgenum, m._impl.flex_elem, @@ -4212,20 +4301,25 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.flex_friction, m._impl.flex_gap, m._impl.flex_internal, + m.flex_interp, m._impl.flex_margin, + m._impl.flex_node, + m.flex_node0, + m.flex_nodeadr, + m.flex_nodebodyid, + m.flex_nodenum, m._impl.flex_priority, m._impl.flex_radius, m._impl.flex_selfcollide, m._impl.flex_shell, - m._impl.flex_shelladr, m._impl.flex_shelldataadr, - m._impl.flex_shellflexid, m._impl.flex_solimp, m._impl.flex_solmix, m._impl.flex_solref, m._impl.flex_stiffness, m._impl.flex_stiffnessadr, m._impl.flex_vert, + m.flex_vert0, m.flex_vertadr, m._impl.flex_vertbodyid, m._impl.flex_vertflexid, @@ -4236,7 +4330,9 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.flexedge_invweight0, m._impl.flexedge_length0, m._impl.flexelem_geom_pair_filtered, - m._impl.flexshell_geom_pair_filtered, + m._impl.flexstrain_J_colind, + m._impl.flexstrain_J_rowadr, + m._impl.flexstrain_J_rownnz, m._impl.flexvert_geom_pair_filtered, m.geom_aabb, m.geom_bodyid, @@ -4262,6 +4358,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.geom_solmix, m.geom_solref, m.geom_type, + m._impl.has_3d_flex, + m._impl.has_ellipsoid_geom, m._impl.has_flex_selfcollide, m._impl.has_fluid, m._impl.has_sdf_geom, @@ -4336,6 +4434,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.nflexedge, m._impl.nflexelem, m._impl.nflexevpair, + m._impl.nflexintcell, + m._impl.nflexnode, m._impl.nflexvert, m.ngeom, m.nhistory, @@ -4545,6 +4645,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.flexedge_J, d._impl.flexedge_length, d._impl.flexedge_velocity, + d._impl.flexnode_xpos, d._impl.flexvert_xpos, d.geom_xmat, d.geom_xpos, @@ -4701,118 +4802,119 @@ def _step_jax_impl(m: types.Model, d: types.Data): '_impl.flexedge_J': out[35], '_impl.flexedge_length': out[36], '_impl.flexedge_velocity': out[37], - '_impl.flexvert_xpos': out[38], - 'geom_xmat': out[39], - 'geom_xpos': out[40], - 'history': out[41], - '_impl.island_dofadr': out[42], - '_impl.island_idofadr': out[43], - '_impl.island_iefcadr': out[44], - '_impl.island_ne': out[45], - '_impl.island_nefc': out[46], - '_impl.island_nf': out[47], - '_impl.island_nv': out[48], - '_impl.light_xdir': out[49], - '_impl.light_xpos': out[50], - '_impl.map_dof2idof': out[51], - '_impl.map_efc2iefc': out[52], - '_impl.map_idof2dof': out[53], - '_impl.map_iefc2efc': out[54], - '_impl.moment_colind': out[55], - '_impl.moment_rowadr': out[56], - '_impl.moment_rownnz': out[57], - '_impl.nacon': out[58], - '_impl.nbody_awake': out[59], - '_impl.ncdof': out[60], - '_impl.ncollision': out[61], - '_impl.ne': out[62], - '_impl.nefc': out[63], - '_impl.nf': out[64], - '_impl.nidof': out[65], - '_impl.nisland': out[66], - '_impl.nl': out[67], - '_impl.ntree_awake': out[68], - '_impl.nv_awake': out[69], - '_impl.overflow': out[70], - '_impl.qLD': out[71], - '_impl.qLDiagInv': out[72], - '_impl.qLU': out[73], - 'qacc': out[74], - 'qacc_smooth': out[75], - 'qacc_warmstart': out[76], - 'qfrc_actuator': out[77], - 'qfrc_bias': out[78], - 'qfrc_constraint': out[79], - '_impl.qfrc_damper': out[80], - 'qfrc_fluid': out[81], - 'qfrc_gravcomp': out[82], - 'qfrc_passive': out[83], - 'qfrc_smooth': out[84], - '_impl.qfrc_spring': out[85], - 'qpos': out[86], - 'qvel': out[87], - 'sensordata': out[88], - 'site_xmat': out[89], - 'site_xpos': out[90], - '_impl.solver_niter': out[91], - '_impl.subtree_angmom': out[92], - 'subtree_com': out[93], - '_impl.subtree_linvel': out[94], - '_impl.ten_J': out[95], - 'ten_length': out[96], - '_impl.ten_velocity': out[97], - '_impl.ten_wrapadr': out[98], - '_impl.ten_wrapnum': out[99], - 'time': out[100], - '_impl.tree_asleep': out[101], - '_impl.tree_awake': out[102], - '_impl.tree_island': out[103], - '_impl.wrap_obj': out[104], - '_impl.wrap_xpos': out[105], - 'xanchor': out[106], - 'xaxis': out[107], - 'ximat': out[108], - 'xipos': out[109], - 'xmat': out[110], - 'xpos': out[111], - 'xquat': out[112], - '_impl.contact__dim': out[113], - '_impl.contact__dist': out[114], - '_impl.contact__efc_address': out[115], - '_impl.contact__elem': out[116], - '_impl.contact__flex': out[117], - '_impl.contact__frame': out[118], - '_impl.contact__friction': out[119], - '_impl.contact__geom': out[120], - '_impl.contact__geomcollisionid': out[121], - '_impl.contact__includemargin': out[122], - '_impl.contact__pos': out[123], - '_impl.contact__solimp': out[124], - '_impl.contact__solref': out[125], - '_impl.contact__solreffriction': out[126], - '_impl.contact__type': out[127], - '_impl.contact__vert': out[128], - '_impl.contact__worldid': out[129], - '_impl.efc__D': out[130], - '_impl.efc__J': out[131], - '_impl.efc__J_colind': out[132], - '_impl.efc__J_rowadr': out[133], - '_impl.efc__J_rownnz': out[134], - '_impl.efc__Jqvel': out[135], - '_impl.efc__Ma': out[136], - '_impl.efc__aref': out[137], - '_impl.efc__force': out[138], - '_impl.efc__frictionloss': out[139], - '_impl.efc__id': out[140], - '_impl.efc__island': out[141], - '_impl.efc__jtdaj_adr': out[142], - '_impl.efc__jtdaj_nblock': out[143], - '_impl.efc__jtdaj_nrow': out[144], - '_impl.efc__margin': out[145], - '_impl.efc__pos': out[146], - '_impl.efc__state': out[147], - '_impl.efc__type': out[148], - '_impl.efc__vel': out[149], + '_impl.flexnode_xpos': out[38], + '_impl.flexvert_xpos': out[39], + 'geom_xmat': out[40], + 'geom_xpos': out[41], + 'history': out[42], + '_impl.island_dofadr': out[43], + '_impl.island_idofadr': out[44], + '_impl.island_iefcadr': out[45], + '_impl.island_ne': out[46], + '_impl.island_nefc': out[47], + '_impl.island_nf': out[48], + '_impl.island_nv': out[49], + '_impl.light_xdir': out[50], + '_impl.light_xpos': out[51], + '_impl.map_dof2idof': out[52], + '_impl.map_efc2iefc': out[53], + '_impl.map_idof2dof': out[54], + '_impl.map_iefc2efc': out[55], + '_impl.moment_colind': out[56], + '_impl.moment_rowadr': out[57], + '_impl.moment_rownnz': out[58], + '_impl.nacon': out[59], + '_impl.nbody_awake': out[60], + '_impl.ncdof': out[61], + '_impl.ncollision': out[62], + '_impl.ne': out[63], + '_impl.nefc': out[64], + '_impl.nf': out[65], + '_impl.nidof': out[66], + '_impl.nisland': out[67], + '_impl.nl': out[68], + '_impl.ntree_awake': out[69], + '_impl.nv_awake': out[70], + '_impl.overflow': out[71], + '_impl.qLD': out[72], + '_impl.qLDiagInv': out[73], + '_impl.qLU': out[74], + 'qacc': out[75], + 'qacc_smooth': out[76], + 'qacc_warmstart': out[77], + 'qfrc_actuator': out[78], + 'qfrc_bias': out[79], + 'qfrc_constraint': out[80], + '_impl.qfrc_damper': out[81], + 'qfrc_fluid': out[82], + 'qfrc_gravcomp': out[83], + 'qfrc_passive': out[84], + 'qfrc_smooth': out[85], + '_impl.qfrc_spring': out[86], + 'qpos': out[87], + 'qvel': out[88], + 'sensordata': out[89], + 'site_xmat': out[90], + 'site_xpos': out[91], + '_impl.solver_niter': out[92], + '_impl.subtree_angmom': out[93], + 'subtree_com': out[94], + '_impl.subtree_linvel': out[95], + '_impl.ten_J': out[96], + 'ten_length': out[97], + '_impl.ten_velocity': out[98], + '_impl.ten_wrapadr': out[99], + '_impl.ten_wrapnum': out[100], + 'time': out[101], + '_impl.tree_asleep': out[102], + '_impl.tree_awake': out[103], + '_impl.tree_island': out[104], + '_impl.wrap_obj': out[105], + '_impl.wrap_xpos': out[106], + 'xanchor': out[107], + 'xaxis': out[108], + 'ximat': out[109], + 'xipos': out[110], + 'xmat': out[111], + 'xpos': out[112], + 'xquat': out[113], + '_impl.contact__dim': out[114], + '_impl.contact__dist': out[115], + '_impl.contact__efc_address': out[116], + '_impl.contact__elem': out[117], + '_impl.contact__flex': out[118], + '_impl.contact__frame': out[119], + '_impl.contact__friction': out[120], + '_impl.contact__geom': out[121], + '_impl.contact__geomcollisionid': out[122], + '_impl.contact__includemargin': out[123], + '_impl.contact__pos': out[124], + '_impl.contact__solimp': out[125], + '_impl.contact__solref': out[126], + '_impl.contact__solreffriction': out[127], + '_impl.contact__type': out[128], + '_impl.contact__vert': out[129], + '_impl.contact__worldid': out[130], + '_impl.efc__D': out[131], + '_impl.efc__J': out[132], + '_impl.efc__J_colind': out[133], + '_impl.efc__J_rowadr': out[134], + '_impl.efc__J_rownnz': out[135], + '_impl.efc__Jqvel': out[136], + '_impl.efc__Ma': out[137], + '_impl.efc__aref': out[138], + '_impl.efc__force': out[139], + '_impl.efc__frictionloss': out[140], + '_impl.efc__id': out[141], + '_impl.efc__island': out[142], + '_impl.efc__jtdaj_adr': out[143], + '_impl.efc__jtdaj_nblock': out[144], + '_impl.efc__jtdaj_nrow': out[145], + '_impl.efc__margin': out[146], + '_impl.efc__pos': out[147], + '_impl.efc__state': out[148], + '_impl.efc__type': out[149], + '_impl.efc__vel': out[150], }) return d diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 18d3449f..ad11625d 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -117,7 +117,7 @@ class BlockDim: render: render block dimension (render) """ segmented_sort: int = 128 - convex_ccd: int = 256 + convex_ccd: int = 64 actuator_velocity: int = 32 ray: int = 64 contact_sort: int = 64 @@ -206,11 +206,14 @@ class ModelWarp(PyTreeNode): dof_tri_row: np.ndarray eq_connect_adr: np.ndarray eq_flex_adr: np.ndarray + eq_flexstrain_adr: np.ndarray eq_jnt_adr: np.ndarray eq_ten_adr: np.ndarray eq_wld_adr: np.ndarray flex_bending: np.ndarray flex_bendingadr: np.ndarray + flex_cell_map: np.ndarray + flex_cellnum: np.ndarray flex_centered: np.ndarray flex_conaffinity: np.ndarray flex_condim: np.ndarray @@ -219,6 +222,7 @@ class ModelWarp(PyTreeNode): flex_dim: np.ndarray flex_edge: np.ndarray flex_edgeadr: np.ndarray + flex_edgeequality: np.ndarray flex_edgeflap: np.ndarray flex_edgenum: np.ndarray flex_elem: np.ndarray @@ -236,6 +240,7 @@ class ModelWarp(PyTreeNode): flex_gap: np.ndarray flex_internal: np.ndarray flex_margin: np.ndarray + flex_node: np.ndarray flex_priority: np.ndarray flex_radius: np.ndarray flex_selfcollide: np.ndarray @@ -258,10 +263,14 @@ class ModelWarp(PyTreeNode): flexedge_invweight0: np.ndarray flexedge_length0: np.ndarray flexelem_geom_pair_filtered: np.ndarray - flexshell_geom_pair_filtered: np.ndarray + flexstrain_J_colind: np.ndarray + flexstrain_J_rowadr: np.ndarray + flexstrain_J_rownnz: np.ndarray flexvert_geom_pair_filtered: np.ndarray geom_pair_type_count: Tuple[int, ...] geom_plugin_index: np.ndarray + has_3d_flex: bool + has_ellipsoid_geom: bool has_flex_selfcollide: bool has_fluid: bool has_sdf_geom: bool @@ -288,14 +297,18 @@ class ModelWarp(PyTreeNode): mesh_polyvertnum: np.ndarray mocap_bodyid: np.ndarray nJfe: int + nJfs: int nacttrnbody: int nbranch: int + neq_flexstrain: int nflexbending: int nflexedge: int nflexelem: int nflexelemdata: int nflexelemedge: int nflexevpair: int + nflexintcell: int + nflexnode: int nflexshelldata: int nflexstiffness: int nflexvert: int @@ -452,6 +465,7 @@ class DataWarp(PyTreeNode): flexedge_J: jax.Array flexedge_length: jax.Array flexedge_velocity: jax.Array + flexnode_xpos: jax.Array flexvert_xpos: jax.Array island_dofadr: jax.Array island_idofadr: jax.Array @@ -664,6 +678,7 @@ _NDIM = { 'flexedge_J': 2, 'flexedge_length': 2, 'flexedge_velocity': 2, + 'flexnode_xpos': 3, 'flexvert_xpos': 3, 'geom_xmat': 4, 'geom_xpos': 3, @@ -882,6 +897,7 @@ _NDIM = { 'eq_connect_adr': 1, 'eq_data': 3, 'eq_flex_adr': 1, + 'eq_flexstrain_adr': 1, 'eq_jnt_adr': 1, 'eq_obj1id': 1, 'eq_obj2id': 1, @@ -894,6 +910,8 @@ _NDIM = { 'exclude_signature': 1, 'flex_bending': 1, 'flex_bendingadr': 1, + 'flex_cell_map': 2, + 'flex_cellnum': 2, 'flex_centered': 1, 'flex_conaffinity': 1, 'flex_condim': 1, @@ -902,6 +920,7 @@ _NDIM = { 'flex_dim': 1, 'flex_edge': 2, 'flex_edgeadr': 1, + 'flex_edgeequality': 1, 'flex_edgeflap': 2, 'flex_edgenum': 1, 'flex_elem': 1, @@ -918,7 +937,13 @@ _NDIM = { 'flex_friction': 2, 'flex_gap': 1, 'flex_internal': 1, + 'flex_interp': 1, 'flex_margin': 1, + 'flex_node': 2, + 'flex_node0': 2, + 'flex_nodeadr': 1, + 'flex_nodebodyid': 1, + 'flex_nodenum': 1, 'flex_priority': 1, 'flex_radius': 1, 'flex_selfcollide': 1, @@ -933,6 +958,7 @@ _NDIM = { 'flex_stiffness': 1, 'flex_stiffnessadr': 1, 'flex_vert': 2, + 'flex_vert0': 2, 'flex_vertadr': 1, 'flex_vertbodyid': 1, 'flex_vertflexid': 1, @@ -943,7 +969,9 @@ _NDIM = { 'flexedge_invweight0': 1, 'flexedge_length0': 1, 'flexelem_geom_pair_filtered': 2, - 'flexshell_geom_pair_filtered': 2, + 'flexstrain_J_colind': 1, + 'flexstrain_J_rowadr': 1, + 'flexstrain_J_rownnz': 1, 'flexvert_geom_pair_filtered': 2, 'geom_aabb': 4, 'geom_bodyid': 1, @@ -969,6 +997,8 @@ _NDIM = { 'geom_solmix': 2, 'geom_solref': 3, 'geom_type': 1, + 'has_3d_flex': 0, + 'has_ellipsoid_geom': 0, 'has_flex_selfcollide': 0, 'has_fluid': 0, 'has_sdf_geom': 0, @@ -1051,6 +1081,7 @@ _NDIM = { 'nC': 0, 'nD': 0, 'nJfe': 0, + 'nJfs': 0, 'nJmom': 0, 'nJten': 0, 'nM': 0, @@ -1060,6 +1091,7 @@ _NDIM = { 'nbranch': 0, 'ncam': 0, 'neq': 0, + 'neq_flexstrain': 0, 'nexclude': 0, 'nflex': 0, 'nflexbending': 0, @@ -1068,6 +1100,8 @@ _NDIM = { 'nflexelemdata': 0, 'nflexelemedge': 0, 'nflexevpair': 0, + 'nflexintcell': 0, + 'nflexnode': 0, 'nflexshelldata': 0, 'nflexstiffness': 0, 'nflexvert': 0, @@ -1365,6 +1399,7 @@ _BATCH_DIM = { 'flexedge_J': True, 'flexedge_length': True, 'flexedge_velocity': True, + 'flexnode_xpos': True, 'flexvert_xpos': True, 'geom_xmat': True, 'geom_xpos': True, @@ -1583,6 +1618,7 @@ _BATCH_DIM = { 'eq_connect_adr': False, 'eq_data': True, 'eq_flex_adr': False, + 'eq_flexstrain_adr': False, 'eq_jnt_adr': False, 'eq_obj1id': False, 'eq_obj2id': False, @@ -1595,6 +1631,8 @@ _BATCH_DIM = { 'exclude_signature': False, 'flex_bending': False, 'flex_bendingadr': False, + 'flex_cell_map': False, + 'flex_cellnum': False, 'flex_centered': False, 'flex_conaffinity': False, 'flex_condim': False, @@ -1603,6 +1641,7 @@ _BATCH_DIM = { 'flex_dim': False, 'flex_edge': False, 'flex_edgeadr': False, + 'flex_edgeequality': False, 'flex_edgeflap': False, 'flex_edgenum': False, 'flex_elem': False, @@ -1619,7 +1658,13 @@ _BATCH_DIM = { 'flex_friction': False, 'flex_gap': False, 'flex_internal': False, + 'flex_interp': False, 'flex_margin': False, + 'flex_node': False, + 'flex_node0': False, + 'flex_nodeadr': False, + 'flex_nodebodyid': False, + 'flex_nodenum': False, 'flex_priority': False, 'flex_radius': False, 'flex_selfcollide': False, @@ -1634,6 +1679,7 @@ _BATCH_DIM = { 'flex_stiffness': False, 'flex_stiffnessadr': False, 'flex_vert': False, + 'flex_vert0': False, 'flex_vertadr': False, 'flex_vertbodyid': False, 'flex_vertflexid': False, @@ -1644,7 +1690,9 @@ _BATCH_DIM = { 'flexedge_invweight0': False, 'flexedge_length0': False, 'flexelem_geom_pair_filtered': False, - 'flexshell_geom_pair_filtered': False, + 'flexstrain_J_colind': False, + 'flexstrain_J_rowadr': False, + 'flexstrain_J_rownnz': False, 'flexvert_geom_pair_filtered': False, 'geom_aabb': True, 'geom_bodyid': False, @@ -1670,6 +1718,8 @@ _BATCH_DIM = { 'geom_solmix': True, 'geom_solref': True, 'geom_type': False, + 'has_3d_flex': False, + 'has_ellipsoid_geom': False, 'has_flex_selfcollide': False, 'has_fluid': False, 'has_sdf_geom': False, @@ -1752,6 +1802,7 @@ _BATCH_DIM = { 'nC': False, 'nD': False, 'nJfe': False, + 'nJfs': False, 'nJmom': False, 'nJten': False, 'nM': False, @@ -1761,6 +1812,7 @@ _BATCH_DIM = { 'nbranch': False, 'ncam': False, 'neq': False, + 'neq_flexstrain': False, 'nexclude': False, 'nflex': False, 'nflexbending': False, @@ -1769,6 +1821,8 @@ _BATCH_DIM = { 'nflexelemdata': False, 'nflexelemedge': False, 'nflexevpair': False, + 'nflexintcell': False, + 'nflexnode': False, 'nflexshelldata': False, 'nflexstiffness': False, 'nflexvert': False,