chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Margin and gap redesign (breaking change)

Browse Source 
PiperOrigin-RevId: 914329812
Change-Id: I905665e4c1965bdb8b90587e5b1277cfbfe5cce0
This commit is contained in:
Yuval Tassa
2026-05-12 09:43:07 -07:00
committed by Copybara-Service
parent 7f5be412a0
commit a4e49f2dff
29 changed files with 2082 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_driver.py
+2 -3
View File
@@ -275,7 +275,7 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
if ip.size > 0:
# pair contacts get their params from m.pair_* fields
params.append((
m.pair_margin[ip] - m.pair_gap[ip],
m.pair_margin[ip],
jp.clip(m.pair_friction[ip], min=eps),
m.pair_solref[ip],
m.pair_solreffriction[ip],
@@ -284,7 +284,6 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
if geom1.size > 0 and geom2.size > 0:
# other contacts get their params from geom fields
margin = m.geom_margin[geom1] + m.geom_margin[geom2]
gap = m.geom_gap[geom1] + m.geom_gap[geom2]
solmix1, solmix2 = m.geom_solmix[geom1], m.geom_solmix[geom2]
mix = solmix1 / (solmix1 + solmix2)
mix = jp.where((solmix1 < eps) & (solmix2 < eps), 0.5, mix)
@@ -315,7 +314,7 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
# unpack 5d friction:
friction = friction[:, [0, 0, 1, 2, 2]]
params.append((margin - gap, friction, solref, solreffriction, solimp))
params.append((margin, friction, solref, solreffriction, solimp))
params = map(jp.concatenate, zip(*params))
includemargin, friction, solref, solreffriction, solimp = params
mjx/mujoco/mjx/_src/types.py
+1 -1
View File
@@ -942,7 +942,7 @@ class Contact(PyTreeNode):
dist: distance between nearest points; neg: penetration
pos: position of contact point: midpoint between geoms (3,)
frame: normal is in [0-2] (9,)
includemargin: include if dist<includemargin=margin-gap (1,)
includemargin: include if dist<includemargin=margin (1,)
friction: tangent1, 2, spin, roll1, 2 (5,)
solref: constraint solver reference, normal direction (mjNREF,)
solreffriction: constraint solver reference, friction directions (mjNREF,)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+11 -3
View File
@@ -35,6 +35,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -773,7 +774,10 @@ def ccd_kernel_builder(
if is_collision_sensor:
cutoff = 1.0e32
else:
cutoff = 0.0
if wp.static(_NEW_GAP_SEMANTICS):
cutoff = gap
else:
cutoff = 0.0
dist, ncollision, w1, w2, multiccd_idx = ccd(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
cutoff,
@@ -793,8 +797,12 @@ def ccd_kernel_builder(
epa_horizon_in[ccdid],
)
if dist >= 0.0 and pairid[1] == -1:
return 0
if wp.static(_NEW_GAP_SEMANTICS):
if dist >= gap and pairid[1] == -1:
return 0
else:
if dist >= 0.0 and pairid[1] == -1:
return 0
# CCD operates on margin-inflated shapes (support() inflates each geom by
# 0.5 * margin). The returned dist is therefore relative to the inflated
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+8 -3
View File
@@ -22,6 +22,7 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -197,14 +198,15 @@ def write_contact(
Returns 1 if the contact is active (dist < margin), 0 otherwise.
"""
active = dist_in < margin_in
detected = dist_in < margin_in + gap_in
# skip contact and no collision sensor
if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1:
if (pairid_in[0] == -2 or not detected) and pairid_in[1] == -1:
return 0
contact_type = 0
if pairid_in[0] >= -1 and active:
if pairid_in[0] >= -1 and detected:
contact_type |= ContactType.CONSTRAINT
if pairid_in[1] >= 0:
@@ -217,7 +219,10 @@ def write_contact(
contact_frame_out[cid] = frame_in
contact_geom_out[cid] = geoms_in
contact_worldid_out[cid] = worldid_in
includemargin = margin_in - gap_in
if wp.static(_NEW_GAP_SEMANTICS):
includemargin = margin_in
else:
includemargin = margin_in - gap_in
contact_includemargin_out[cid] = includemargin
contact_dim_out[cid] = condim_in
contact_friction_out[cid] = friction_in
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+25 -15
View File
@@ -271,13 +271,14 @@ def _obb_filter(
return True
def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
@wp.func
def func(
# Model:
geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d[float],
geom_margin: wp.array2d[float],
geom_gap: wp.array2d[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
@@ -299,21 +300,25 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2] # kernel_analyzer: ignore
margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] # kernel_analyzer: ignore
gap_id = worldid % ngeom_gap if wp.static(ngeom_gap > 1) else 0
gap1, gap2 = geom_gap[gap_id, geom1], geom_gap[gap_id, geom2] # kernel_analyzer: ignore
effective_margin1 = margin1 + gap1
effective_margin2 = margin2 + gap2
xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
if rbound1 == 0.0 or rbound2 == 0.0:
if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE):
return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
return _plane_filter(rbound1, rbound2, effective_margin1, effective_margin2, xpos1, xpos2, xmat1, xmat2)
else:
if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE):
if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
if not _sphere_filter(rbound1, rbound2, effective_margin1, effective_margin2, xpos1, xpos2):
return False
if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB):
if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
if not _aabb_filter(center1, center2, size1, size2, effective_margin1, effective_margin2, xpos1, xpos2, xmat1, xmat2):
return False
if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB):
if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
if not _obb_filter(center1, center2, size1, size2, effective_margin1, effective_margin2, xpos1, xpos2, xmat1, xmat2):
return False
return True
@@ -377,6 +382,7 @@ def _sap_project(opt_broadphase: int):
ngeom: int,
geom_rbound: wp.array2d[float],
geom_margin: wp.array2d[float],
geom_gap: wp.array2d[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
nworld_in: int,
@@ -397,7 +403,7 @@ def _sap_project(opt_broadphase: int):
# geom is a plane
rbound = MJ_MAXVAL
radius = rbound + geom_margin[worldid % geom_margin.shape[0], geomid]
radius = rbound + geom_margin[worldid % geom_margin.shape[0], geomid] + geom_gap[worldid % geom_gap.shape[0], geomid]
center = wp.dot(direction_in, xpos)
sort_index_out[worldid, geomid] = geomid
@@ -443,7 +449,7 @@ def _sap_range(
@cache_kernel
def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
@@ -452,6 +458,7 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d[float],
geom_margin: wp.array2d[float],
geom_gap: wp.array2d[float],
nxn_pairid: wp.array[wp.vec2i],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
@@ -502,8 +509,8 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
continue
if (
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
)
or pairid[1] >= 0
):
@@ -586,7 +593,7 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
wp.launch(
kernel=_sap_project(m.opt.broadphase),
dim=(d.nworld, m.ngeom),
inputs=[m.ngeom, m.geom_rbound, m.geom_margin, d.geom_xpos, d.nworld, direction],
inputs=[m.ngeom, m.geom_rbound, m.geom_margin, m.geom_gap, d.geom_xpos, d.nworld, direction],
outputs=[
projection_lower.reshape((-1, m.ngeom)),
projection_upper,
@@ -622,7 +629,7 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
# assumes each geom has 5 other geoms (batched over all worlds)
nsweep = 5 * nworldgeom
wp.launch(
kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
dim=nsweep,
inputs=[
m.ngeom,
@@ -630,6 +637,7 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
m.geom_aabb,
m.geom_rbound,
m.geom_margin,
m.geom_gap,
m.nxn_pairid,
d.geom_xpos,
d.geom_xmat,
@@ -644,7 +652,7 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
@cache_kernel
def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
@@ -652,6 +660,7 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d[float],
geom_margin: wp.array2d[float],
geom_gap: wp.array2d[float],
nxn_geom_pair: wp.array[wp.vec2i],
nxn_pairid: wp.array[wp.vec2i],
# Data in:
@@ -672,8 +681,8 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
geom2 = geom[1]
if (
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
)
or nxn_pairid[elementid][1] >= 0
):
@@ -709,13 +718,14 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
"""
wp.launch(
_nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
_nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[
m.geom_type,
m.geom_aabb,
m.geom_rbound,
m.geom_margin,
m.geom_gap,
m.nxn_geom_pair_filtered,
m.nxn_pairid_filtered,
d.geom_xpos,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+6 -3
View File
@@ -17,6 +17,7 @@ import enum
from typing import Callable
import mujoco
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
import numpy as np
import warp as wp
@@ -26,6 +27,8 @@ MJ_MINIMP = mujoco.mjMINIMP # minimum constraint impedance
MJ_MAXIMP = mujoco.mjMAXIMP # maximum constraint impedance
MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
MJ_MINMU = mujoco.mjMINMU # minimum friction
# True if MuJoCo >= 3.9.0 (new margin/gap semantics: includemargin = margin)
_NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0")
# maximum size (by number of edges) of an horizon in EPA algorithm
MJ_MAX_EPAHORIZON = 24
# maximum average number of trianglarfaces EPA can insert at each iteration
@@ -974,7 +977,7 @@ class Model:
geom_quat: local orientation offset rel. to body (*, ngeom, 4)
geom_friction: friction for (slide, spin, roll) (*, ngeom, 3)
geom_margin: detect contact if dist<margin (*, ngeom,)
geom_gap: include in solver if dist<margin-gap (*, ngeom,)
geom_gap: additional contact detection buffer (*, ngeom,)
geom_fluid: fluid interaction parameters (ngeom, mjNFLUID)
geom_rgba: rgba when material is omitted (*, ngeom, 4)
site_type: geom type for rendering (GeomType) (nsite,)
@@ -1082,7 +1085,7 @@ class Model:
pair_solreffriction: solver reference: contact friction (*, npair, mjNREF)
pair_solimp: solver impedance: contact (*, npair, mjNIMP)
pair_margin: detect contact if dist<margin (*, npair,)
pair_gap: include in solver if dist<margin-gap (*, npair,)
pair_gap: additional contact detection buffer (*, npair,)
pair_friction: tangent1, 2, spin, roll1, 2 (*, npair, 5)
exclude_signature: body1 << 16 + body2 (nexclude,)
eq_type: constraint type (EqType) (neq,)
@@ -1663,7 +1666,7 @@ class Contact:
dist: distance between nearest points; neg: penetration (naconmax,)
pos: position of contact point: midpoint between geoms (naconmax, 3)
frame: normal is in [0-2], points from geom[0] to geom[1] (naconmax, 3, 3)
includemargin: include if dist<includemargin=margin-gap (naconmax,)
includemargin: include if dist<includemargin=margin (naconmax,)
friction: tangent1, 2, spin, roll1, 2 (naconmax, 5)
solref: constraint solver reference, normal direction (naconmax, 2)
solreffriction: constraint solver reference, friction directions (naconmax, 2)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
Vendored
+12 -8
View File
@@ -22,21 +22,25 @@ import re
def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
"""Parse a version string into comparable components.
Both '.' and '-' are treated as separators. Each component is wrapped in a
tuple: (0, int) for numeric parts, (-1, str) for non-numeric. A (0, 0)
sentinel is appended so that stable releases sort above pre-release suffixes
during Python tuple comparison (e.g., 1.2.3 >= 1.2.3.dev). Non-numeric
components are compared lexicographically (e.g., b >= a).
Dot-separated components form the version. Hyphen-separated suffixes (e.g.,
"-foo3") are treated as local build identifiers and stripped before
parsing. Each component is wrapped in a tuple: (0, int) for numeric parts,
(-1, str) for non-numeric. A (0, 0) sentinel is appended so that stable
releases sort above pre-release suffixes during Python tuple comparison
(e.g., 1.2.3 >= 1.2.3.dev). Non-numeric components are compared
lexicographically (e.g., b >= a).
Args:
version_str: Version string like "3.5.0" or "3.5.0.dev869102767".
version_str: Version string like "3.5.0", "3.5.0.dev869102767", or
"3.9.0-foo3".
Returns:
Tuple of (type_order, value) pairs for comparison, where type_order is 0
for integers and -1 for strings, followed by a (0, 0) sentinel.
"""
# Split on both '.' and '-'
parts = re.split(r"[.\-]", version_str)
# Strip local build identifier (e.g., "3.9.0-foo3" -> "3.9.0")
version_str = version_str.split("-", 1)[0]
parts = version_str.split(".")
return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])