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,