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

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
#MIGRATION_3P_PY_MUJOCO_WARP__DEFAULT

  - 19dde09e179dbbaf4d07084185074f6a774b3255 update derivative.py by Taylor Howell <taylorhowell@google.com>
  - 04bc115c3766e974ef000cca16d40bacb92c01d2 fix _ray_hfield by Taylor Howell <taylorhowell@google.com>
  - 7fcbf8ddba08c9b9d33b0fdd4dd3fe39e40a0bd7 Use sparsity structures in `mjModel` rather than `mjData` by Yuval Tassa <6387526+yuvaltassa@users.noreply.github.com>
  - bfd2e4fe3cb5ba5dee97445ddc1d072344f443ec Fix override for enums (#617) by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - f4dbf5345bd5dcae60c415361dd425605b672e5c Add keyframe selection to viewer (#616) by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - 0e7edf9e8fd337502586929b2e4dd3d51201f462 add prunning to multiccd by Kyle Bayes <42924292+kbayes@users.noreply.github.com>
  - 791c30237cbdecf876bc53ac0b3dcff8b8b1317d format by Kyle Bayes <42924292+kbayes@users.noreply.github.com>
  - 7a7bae6c1390df3f29058e90f25f9ed5caab3fb1 Fix override in viewer and testspeed (#620) by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - 9e230c229c7dd721fbf6fa5fe69fc98717895500 add + 0.0 by Taylor Howell <taylorhowell@google.com>
  - 8a70a5e8839b84c11601bf1a98f9e7d1d31603f1 Update the test_cost unit test to test more properties (#... by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - cf317c60845b24ddba4b6814dbac9eba32b4c8b0 parallel linesearch with log scale for step sizes (#575) by Taylor Howell <taylorhowell@google.com>
  - b73e4edbc0d331d1e4f15c006d4535cc59595e70 add better reporting (#600) by adenzler-nvidia <116633880+adenzler-nvidia@users.noreply.github.com>
  - 007ad0b204ed4d48584024cce4676ac25a5695b5 Optimization: Fusing kernels around update_gradient_JTDAJ... by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - 640a89fcf452b9564508cacdeb30ecef3a83aede Adding static to broadphase filter (#621) by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - daca09ff3261db2713abc912191d6cf76437b8b7 Add reporting of converged worlds in testspeed (#619) by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - 33830bddd95a7f0d3febe1b090398ddadeb86ee2 Fix call to benchmark from asv (#627) by Kenny-Vilella <163487060+Kenny-Vilella@users.noreply.github.com>
  - 382eb65f7bb45798698f4025a980287ab1ec84c7 Update flex_bending dimensions (#628) by Alessio Quaglino <102665781+quagla@users.noreply.github.com>
  - c56b8875439cfaa93c1393f0a6e2ad057e029064 contact sensor maxforce reduction (#602) by Taylor Howell <taylorhowell@google.com>
  - 4f731941cb9dea397d4f71cf8151100b83eae576 A few airspeed velocity improvements. (#626) by Erik Frey <misc@fawx.com>

PiperOrigin-RevId: 797054114
Change-Id: I380e0b9446f94202b2389e2a760d6de5f0ec5af7
This commit is contained in:
Baruch Tabanpour
2025-08-19 16:00:39 -07:00
committed by Copybara-Service
parent fd96de330e
commit bf04103b6f
21 changed files with 528 additions and 430 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+162 -158
View File
@@ -29,7 +29,9 @@ 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 DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
wp.set_module_options({"enable_backward": False}) wp.set_module_options({"enable_backward": False})
@@ -248,54 +250,56 @@ def _obb_filter(
return True return True
@wp.func def _broadphase_filter(opt_broadphase_filter: int):
def _broadphase_filter( @wp.func
# Model: def func(
opt_broadphase_filter: int, # Model:
geom_aabb: wp.array2d(dtype=wp.vec3), geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d(dtype=float),
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In: # In:
geom1: int, geom1: int,
geom2: int, geom2: int,
worldid: int, worldid: int,
) -> bool: ) -> bool:
# 1: plane # 1: plane
# 2: sphere # 2: sphere
# 4: aabb # 4: aabb
# 8: obb # 8: obb
center1 = geom_aabb[geom1, 0] center1 = geom_aabb[geom1, 0]
center2 = geom_aabb[geom2, 0] center2 = geom_aabb[geom2, 0]
size1 = geom_aabb[geom1, 1] size1 = geom_aabb[geom1, 1]
size2 = geom_aabb[geom2, 1] size2 = geom_aabb[geom2, 1]
rbound1 = geom_rbound[worldid, geom1] rbound1 = geom_rbound[worldid, geom1]
rbound2 = geom_rbound[worldid, geom2] rbound2 = geom_rbound[worldid, geom2]
margin1 = geom_margin[worldid, geom1] margin1 = geom_margin[worldid, geom1]
margin2 = geom_margin[worldid, geom2] margin2 = geom_margin[worldid, geom2]
xpos1 = geom_xpos_in[worldid, geom1] xpos1 = geom_xpos_in[worldid, geom1]
xpos2 = geom_xpos_in[worldid, geom2] xpos2 = geom_xpos_in[worldid, geom2]
xmat1 = geom_xmat_in[worldid, geom1] xmat1 = geom_xmat_in[worldid, geom1]
xmat2 = geom_xmat_in[worldid, geom2] xmat2 = geom_xmat_in[worldid, geom2]
if rbound1 == 0.0 or rbound2 == 0.0: if rbound1 == 0.0 or rbound2 == 0.0:
if opt_broadphase_filter & int(BroadphaseFilter.PLANE.value): if wp.static(opt_broadphase_filter & int(BroadphaseFilter.PLANE.value)):
return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2) return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
else: else:
if opt_broadphase_filter & int(BroadphaseFilter.SPHERE.value): if wp.static(opt_broadphase_filter & int(BroadphaseFilter.SPHERE.value)):
if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2): if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
return False return False
if opt_broadphase_filter & int(BroadphaseFilter.AABB.value): if wp.static(opt_broadphase_filter & int(BroadphaseFilter.AABB.value)):
if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False return False
if opt_broadphase_filter & int(BroadphaseFilter.OBB.value): if wp.static(opt_broadphase_filter & int(BroadphaseFilter.OBB.value)):
if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False return False
return True return True
return func
@wp.func @wp.func
@@ -413,82 +417,81 @@ def _sap_range(
sap_range_out[worldid, geomid] = limit - geomid sap_range_out[worldid, geomid] = limit - geomid
@wp.kernel @cache_kernel
def _sap_broadphase( def _sap_broadphase(broadphase_filter):
# Model: @nested_kernel
ngeom: int, def kernel(
opt_broadphase_filter: int, # Model:
geom_type: wp.array(dtype=int), ngeom: int,
geom_aabb: wp.array2d(dtype=wp.vec3), geom_type: wp.array(dtype=int),
geom_rbound: wp.array2d(dtype=float), geom_aabb: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float), geom_rbound: wp.array2d(dtype=float),
nxn_pairid: wp.array(dtype=int), geom_margin: wp.array2d(dtype=float),
# Data in: nxn_pairid: wp.array(dtype=int),
nworld_in: int, # Data in:
nconmax_in: int, nworld_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3), nconmax_in: int,
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3),
sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore geom_xmat_in: wp.array2d(dtype=wp.mat33),
sap_cumulative_sum_in: wp.array(dtype=int), # kernel_analyzer: ignore sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore
# In: sap_cumulative_sum_in: wp.array(dtype=int), # kernel_analyzer: ignore
nsweep_in: int, # In:
# Data out: nsweep_in: int,
collision_pair_out: wp.array(dtype=wp.vec2i), # Data out:
collision_hftri_index_out: wp.array(dtype=int), collision_pair_out: wp.array(dtype=wp.vec2i),
collision_pairid_out: wp.array(dtype=int), collision_hftri_index_out: wp.array(dtype=int),
collision_worldid_out: wp.array(dtype=int), collision_pairid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int), collision_worldid_out: wp.array(dtype=int),
): ncollision_out: wp.array(dtype=int),
worldgeomid = wp.tid() ):
worldgeomid = wp.tid()
nworldgeom = nworld_in * ngeom nworldgeom = nworld_in * ngeom
nworkpackages = sap_cumulative_sum_in[nworldgeom - 1] nworkpackages = sap_cumulative_sum_in[nworldgeom - 1]
while worldgeomid < nworkpackages: while worldgeomid < nworkpackages:
# binary search to find current and next geom pair indices # binary search to find current and next geom pair indices
i = _binary_search(sap_cumulative_sum_in, worldgeomid, 0, nworldgeom) i = _binary_search(sap_cumulative_sum_in, worldgeomid, 0, nworldgeom)
j = i + worldgeomid + 1 j = i + worldgeomid + 1
if i > 0: if i > 0:
j -= sap_cumulative_sum_in[i - 1] j -= sap_cumulative_sum_in[i - 1]
worldid = i // ngeom worldid = i // ngeom
i = i % ngeom i = i % ngeom
j = j % ngeom j = j % ngeom
# get geom indices and swap if necessary # get geom indices and swap if necessary
geom1 = sap_sort_index_in[worldid, i] geom1 = sap_sort_index_in[worldid, i]
geom2 = sap_sort_index_in[worldid, j] geom2 = sap_sort_index_in[worldid, j]
# find linear index of (geom1, geom2) in upper triangular nxn_pairid # find linear index of (geom1, geom2) in upper triangular nxn_pairid
if geom2 < geom1: if geom2 < geom1:
idx = upper_tri_index(ngeom, geom2, geom1) idx = upper_tri_index(ngeom, geom2, geom1)
else: else:
idx = upper_tri_index(ngeom, geom1, geom2) idx = upper_tri_index(ngeom, geom1, geom2)
if nxn_pairid[idx] < -1:
worldgeomid += nsweep_in worldgeomid += nsweep_in
continue if nxn_pairid[idx] < -1:
continue
if _broadphase_filter( if broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid):
opt_broadphase_filter, geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid _add_geom_pair(
): geom_type,
_add_geom_pair( nxn_pairid,
geom_type, nconmax_in,
nxn_pairid, geom1,
nconmax_in, geom2,
geom1, worldid,
geom2, idx,
worldid, collision_pair_out,
idx, collision_hftri_index_out,
collision_pair_out, collision_pairid_out,
collision_hftri_index_out, collision_worldid_out,
collision_pairid_out, ncollision_out,
collision_worldid_out, )
ncollision_out,
)
worldgeomid += nsweep_in return kernel
def _segmented_sort(tile_size: int): def _segmented_sort(tile_size: int):
@@ -591,12 +594,12 @@ def sap_broadphase(m: Model, d: Data):
# estimate number of overlap checks # estimate number of overlap checks
# assumes each geom has 5 other geoms (batched over all worlds) # assumes each geom has 5 other geoms (batched over all worlds)
nsweep = 5 * nworldgeom nsweep = 5 * nworldgeom
broadphase_filter = _broadphase_filter(m.opt.broadphase_filter)
wp.launch( wp.launch(
kernel=_sap_broadphase, kernel=_sap_broadphase(broadphase_filter),
dim=nsweep, dim=nsweep,
inputs=[ inputs=[
m.ngeom, m.ngeom,
m.opt.broadphase_filter,
m.geom_type, m.geom_type,
m.geom_aabb, m.geom_aabb,
m.geom_rbound, m.geom_rbound,
@@ -620,50 +623,51 @@ def sap_broadphase(m: Model, d: Data):
) )
@wp.kernel @cache_kernel
def _nxn_broadphase( def _nxn_broadphase(broadphase_filter):
# Model: @nested_kernel
opt_broadphase_filter: int, def kernel(
geom_type: wp.array(dtype=int), # Model:
geom_aabb: wp.array2d(dtype=wp.vec3), geom_type: wp.array(dtype=int),
geom_rbound: wp.array2d(dtype=float), geom_aabb: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float), geom_rbound: wp.array2d(dtype=float),
nxn_geom_pair: wp.array(dtype=wp.vec2i), geom_margin: wp.array2d(dtype=float),
nxn_pairid: wp.array(dtype=int), nxn_geom_pair: wp.array(dtype=wp.vec2i),
# Data in: nxn_pairid: wp.array(dtype=int),
nconmax_in: int, # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), nconmax_in: int,
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3),
# Data out: geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_out: wp.array(dtype=wp.vec2i), # Data out:
collision_hftri_index_out: wp.array(dtype=int), collision_pair_out: wp.array(dtype=wp.vec2i),
collision_pairid_out: wp.array(dtype=int), collision_hftri_index_out: wp.array(dtype=int),
collision_worldid_out: wp.array(dtype=int), collision_pairid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int), collision_worldid_out: wp.array(dtype=int),
): ncollision_out: wp.array(dtype=int),
worldid, elementid = wp.tid()
geom = nxn_geom_pair[elementid]
geom1 = geom[0]
geom2 = geom[1]
if _broadphase_filter(
opt_broadphase_filter, geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
): ):
_add_geom_pair( worldid, elementid = wp.tid()
geom_type,
nxn_pairid, geom = nxn_geom_pair[elementid]
nconmax_in, geom1 = geom[0]
geom1, geom2 = geom[1]
geom2,
worldid, if broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid):
elementid, _add_geom_pair(
collision_pair_out, geom_type,
collision_hftri_index_out, nxn_pairid,
collision_pairid_out, nconmax_in,
collision_worldid_out, geom1,
ncollision_out, geom2,
) worldid,
elementid,
collision_pair_out,
collision_hftri_index_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
return kernel
@event_scope @event_scope
@@ -681,11 +685,11 @@ def nxn_broadphase(m: Model, d: Data):
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags. `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
""" """
broadphase_filter = _broadphase_filter(m.opt.broadphase_filter)
wp.launch( wp.launch(
_nxn_broadphase, _nxn_broadphase(broadphase_filter),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]), dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[ inputs=[
m.opt.broadphase_filter,
m.geom_type, m.geom_type,
m.geom_aabb, m.geom_aabb,
m.geom_rbound, m.geom_rbound,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+113 -54
View File
@@ -41,6 +41,7 @@ polyindices = wp.types.vector(MAX_POLYVERT, dtype=int)
mat43 = wp.types.matrix(shape=(4, 3), dtype=float) mat43 = wp.types.matrix(shape=(4, 3), dtype=float)
mat63 = wp.types.matrix(shape=(6, 3), dtype=float) mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
# box-box currently supports up to 8 contacts
MULTI_CONTACT_COUNT = 8 MULTI_CONTACT_COUNT = 8
mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float) mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
@@ -94,7 +95,7 @@ class SupportPoint:
@wp.func @wp.func
def discrete_geoms(g1: int, g2: int): def _discrete_geoms(g1: int, g2: int):
return (g1 == int(GeomType.MESH.value) or g1 == int(GeomType.BOX.value) or g1 == int(GeomType.HFIELD.value)) and ( return (g1 == int(GeomType.MESH.value) or g1 == int(GeomType.BOX.value) or g1 == int(GeomType.HFIELD.value)) and (
g2 == int(GeomType.MESH.value) or g2 == int(GeomType.BOX.value) or g2 == int(GeomType.HFIELD.value) g2 == int(GeomType.MESH.value) or g2 == int(GeomType.BOX.value) or g2 == int(GeomType.HFIELD.value)
) )
@@ -597,7 +598,7 @@ def _gjk(
use_margin: bool, use_margin: bool,
): ):
"""Find distance within a tolerance between two geoms.""" """Find distance within a tolerance between two geoms."""
is_discrete = discrete_geoms(geomtype1, geomtype2) is_discrete = _discrete_geoms(geomtype1, geomtype2)
cutoff2 = cutoff * cutoff cutoff2 = cutoff * cutoff
simplex = mat43() simplex = mat43()
simplex1 = mat43() simplex1 = mat43()
@@ -1203,7 +1204,7 @@ def _polytope4(
@wp.func @wp.func
def _epa(tolerance: float, epa_iterations: int, pt: Polytope, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int): def _epa(tolerance: float, epa_iterations: int, pt: Polytope, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int):
"""Recover penetration data from two geoms in contact given an initial polytope.""" """Recover penetration data from two geoms in contact given an initial polytope."""
is_discrete = discrete_geoms(geomtype1, geomtype2) is_discrete = _discrete_geoms(geomtype1, geomtype2)
upper = FLOAT_MAX upper = FLOAT_MAX
upper2 = FLOAT_MAX upper2 = FLOAT_MAX
idx = int(-1) idx = int(-1)
@@ -1309,9 +1310,60 @@ def _epa(tolerance: float, epa_iterations: int, pt: Polytope, geom1: Geom, geom2
return 0.0, wp.vec3(), wp.vec3(), -1 return 0.0, wp.vec3(), wp.vec3(), -1
@wp.func
def _area4(a: wp.vec3, b: wp.vec3, c: wp.vec3, d: wp.vec3):
"""Computes area of a quadrilateral embedded in 3D space."""
return 0.5 * wp.norm_l2(wp.cross(a - d, d - b) + wp.cross(b - c, c - a))
@wp.func
def _next(n: int, i: int):
"""Returns (i + 1) mod n for 0 <= i <= n - 1."""
return wp.where(i == n - 1, 0, i + 1)
@wp.func
def _polygon_quad(polygon: polyclip, npolygon: int):
"""Returns the indices of a quadrilateral of maximum area in a convex polygon."""
b = _next(npolygon, 0)
c = _next(npolygon, b)
d = _next(npolygon, c)
res = wp.vec4i(0, b, c, d)
m = _area4(polygon[0], polygon[b], polygon[c], polygon[d])
for a in range(npolygon):
while True:
m_next = _area4(polygon[a], polygon[b], polygon[c], polygon[_next(npolygon, d)])
if m_next <= m:
break
m = m_next
d = _next(npolygon, d)
res = wp.vec4i(a, b, c, d)
while True:
m_next = _area4(polygon[a], polygon[b], polygon[_next(npolygon, c)], polygon[d])
if m_next <= m:
break
m = m_next
c = _next(npolygon, c)
res = wp.vec4i(a, b, c, d)
while True:
m_next = _area4(polygon[a], polygon[_next(npolygon, b)], polygon[c], polygon[d])
if m_next <= m:
break
m = m_next
b = _next(npolygon, b)
res = wp.vec4i(a, b, c, d)
if b == a:
b = _next(npolygon, b)
if c == b:
c = _next(npolygon, c)
if d == c:
d == _next(npolygon, d)
return res
# return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary # return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary
@wp.func @wp.func
def feature_dim(face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3)): def _feature_dim(face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3)):
v1i = vert_index[face[0]] v1i = vert_index[face[0]]
v2i = vert_index[face[1]] v2i = vert_index[face[1]]
v3i = vert_index[face[2]] v3i = vert_index[face[2]]
@@ -1335,7 +1387,7 @@ def feature_dim(face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(
# find two normals that are facing each other within a tolerance, return 1 if found # find two normals that are facing each other within a tolerance, return 1 if found
@wp.func @wp.func
def aligned_faces(vert1: polyverts, len1: int, vert2: polyverts, len2: int): def _aligned_faces(vert1: polyverts, len1: int, vert2: polyverts, len2: int):
res = wp.vec2i() res = wp.vec2i()
for i in range(len1): for i in range(len1):
for j in range(len2): for j in range(len2):
@@ -1349,7 +1401,7 @@ def aligned_faces(vert1: polyverts, len1: int, vert2: polyverts, len2: int):
# find two normals that are perpendicular to each other within a tolerance # find two normals that are perpendicular to each other within a tolerance
# return 1 if found # return 1 if found
@wp.func @wp.func
def aligned_face_edge(edge: polyverts, nedge: int, face: polyverts, nface: int): def _aligned_face_edge(edge: polyverts, nedge: int, face: polyverts, nface: int):
res = wp.vec2i() res = wp.vec2i()
for i in range(nface): for i in range(nface):
for j in range(nedge): for j in range(nedge):
@@ -1362,7 +1414,7 @@ def aligned_face_edge(edge: polyverts, nedge: int, face: polyverts, nface: int):
# find up to n <= 2 common integers of two arrays, return n # find up to n <= 2 common integers of two arrays, return n
@wp.func @wp.func
def intersect1(a1: wp.array(dtype=int), a2: wp.array(dtype=int), start1: int, start2: int, len1: int, len2: int): def _intersect1(a1: wp.array(dtype=int), a2: wp.array(dtype=int), start1: int, start2: int, len1: int, len2: int):
count = int(0) count = int(0)
res = wp.vec2i() res = wp.vec2i()
for i in range(start1, start1 + len1): for i in range(start1, start1 + len1):
@@ -1376,7 +1428,7 @@ def intersect1(a1: wp.array(dtype=int), a2: wp.array(dtype=int), start1: int, st
@wp.func @wp.func
def intersect2(a1: wp.vec2i, a2: wp.array(dtype=int), start2: int, len1: int, len2: int): def _intersect2(a1: wp.vec2i, a2: wp.array(dtype=int), start2: int, len1: int, len2: int):
count = int(0) count = int(0)
res = wp.vec2i() res = wp.vec2i()
for i in range(len1): for i in range(len1):
@@ -1391,7 +1443,7 @@ def intersect2(a1: wp.vec2i, a2: wp.array(dtype=int), start2: int, len1: int, le
# compute possible polygon normals of a mesh given up to 3 vertices # compute possible polygon normals of a mesh given up to 3 vertices
@wp.func @wp.func
def mesh_normals( def _mesh_normals(
# In: # In:
feature_dim: int, feature_dim: int,
feature_index: wp.vec3i, feature_index: wp.vec3i,
@@ -1420,10 +1472,10 @@ def mesh_normals(
v3_num = polymapnum[vertadr + v3] v3_num = polymapnum[vertadr + v3]
faceset = wp.vec2i() faceset = wp.vec2i()
n, edgeset = intersect1(polymap, polymap, v1_adr, v2_adr, v1_num, v2_num) n, edgeset = _intersect1(polymap, polymap, v1_adr, v2_adr, v1_num, v2_num)
if n == 0: if n == 0:
return 0, normals, indices return 0, normals, indices
n, faceset = intersect2(edgeset, polymap, v3_adr, n, v3_num) n, faceset = _intersect2(edgeset, polymap, v3_adr, n, v3_num)
if n == 0: if n == 0:
return 0, normals, indices return 0, normals, indices
@@ -1440,7 +1492,7 @@ def mesh_normals(
v2_num = polymapnum[vertadr + v2] v2_num = polymapnum[vertadr + v2]
# up to two faces as two vertices define an edge # up to two faces as two vertices define an edge
n, edgeset = intersect1(polymap, polymap, v1_adr, v2_adr, v1_num, v2_num) n, edgeset = _intersect1(polymap, polymap, v1_adr, v2_adr, v1_num, v2_num)
if n == 0: if n == 0:
return 0, normals, indices return 0, normals, indices
for i in range(n): for i in range(n):
@@ -1462,7 +1514,7 @@ def mesh_normals(
# compute normal directional vectors along possible edges given by up to two vertices # compute normal directional vectors along possible edges given by up to two vertices
@wp.func @wp.func
def mesh_edge_normals( def _mesh_edge_normals(
# In: # In:
dim: int, dim: int,
mat: wp.mat33, mat: wp.mat33,
@@ -1511,7 +1563,7 @@ def mesh_edge_normals(
# try recovering box normal from collision normal # try recovering box normal from collision normal
@wp.func @wp.func
def box_normals2(mat: wp.mat33, n: wp.vec3): def _box_normals2(mat: wp.mat33, n: wp.vec3):
normals = polyverts() normals = polyverts()
indices = polyindices() indices = polyindices()
@@ -1540,7 +1592,7 @@ def box_normals2(mat: wp.mat33, n: wp.vec3):
# compute possible face normals of a box given up to 3 vertices # compute possible face normals of a box given up to 3 vertices
@wp.func @wp.func
def box_normals(feature_dim: int, feature_index: wp.vec3i, mat: wp.mat33, dir: wp.vec3): def _box_normals(feature_dim: int, feature_index: wp.vec3i, mat: wp.mat33, dir: wp.vec3):
normals = polyverts() normals = polyverts()
indices = polyindices() indices = polyindices()
@@ -1568,7 +1620,7 @@ def box_normals(feature_dim: int, feature_index: wp.vec3i, mat: wp.mat33, dir: w
indices[0] = indices[0] + 1 indices[0] = indices[0] + 1
if c == 1: if c == 1:
return 1, normals, indices return 1, normals, indices
return box_normals2(mat, dir) return _box_normals2(mat, dir)
if feature_dim == 2: if feature_dim == 2:
c = 0 c = 0
x = float((v1 & 1) and (v2 & 1)) - float(not (v1 & 1) and not (v2 & 1)) x = float((v1 & 1) and (v2 & 1)) - float(not (v1 & 1) and not (v2 & 1))
@@ -1588,7 +1640,7 @@ def box_normals(feature_dim: int, feature_index: wp.vec3i, mat: wp.mat33, dir: w
c += 1 c += 1
if c == 2: if c == 2:
return 2, normals, indices return 2, normals, indices
return box_normals2(mat, dir) return _box_normals2(mat, dir)
if feature_dim == 1: if feature_dim == 1:
x = wp.where(v1 & 1, 1.0, -1.0) x = wp.where(v1 & 1, 1.0, -1.0)
@@ -1606,7 +1658,7 @@ def box_normals(feature_dim: int, feature_index: wp.vec3i, mat: wp.mat33, dir: w
# compute possible edge normals for box for edge collisions # compute possible edge normals for box for edge collisions
@wp.func @wp.func
def box_edge_normals(dim: int, mat: wp.mat33, pos: wp.vec3, size: wp.vec3, v1: wp.vec3, v2: wp.vec3, v1i: int): def _box_edge_normals(dim: int, mat: wp.mat33, pos: wp.vec3, size: wp.vec3, v1: wp.vec3, v2: wp.vec3, v1i: int):
normals = polyverts() normals = polyverts()
endverts = polyverts() endverts = polyverts()
@@ -1635,7 +1687,7 @@ def box_edge_normals(dim: int, mat: wp.mat33, pos: wp.vec3, size: wp.vec3, v1: w
# recover face of a box from its index # recover face of a box from its index
@wp.func @wp.func
def box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int): def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int):
res = polyverts() res = polyverts()
# compute global coordinates of the box face and face normal # compute global coordinates of the box face and face normal
@@ -1680,7 +1732,7 @@ def box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int):
# recover mesh polygon from its index, return number of edges # recover mesh polygon from its index, return number of edges
@wp.func @wp.func
def mesh_face( def _mesh_face(
# In: # In:
mat: wp.mat33, mat: wp.mat33,
pos: wp.vec3, pos: wp.vec3,
@@ -1706,19 +1758,19 @@ def mesh_face(
@wp.func @wp.func
def plane_normal(v1: wp.vec3, v2: wp.vec3, n: wp.vec3): def _plane_normal(v1: wp.vec3, v2: wp.vec3, n: wp.vec3):
v3 = v1 + n v3 = v1 + n
res = wp.cross(v2 - v1, v3 - v1) res = wp.cross(v2 - v1, v3 - v1)
return wp.dot(res, v1), res return wp.dot(res, v1), res
@wp.func @wp.func
def halfspace(a: wp.vec3, n: wp.vec3, p: wp.vec3): def _halfspace(a: wp.vec3, n: wp.vec3, p: wp.vec3):
return wp.dot(p - a, n) > -1e-10 return wp.dot(p - a, n) > -1e-10
@wp.func @wp.func
def plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3): def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3):
res = wp.vec3() res = wp.vec3()
ab = b - a ab = b - a
temp = wp.dot(pn, ab) temp = wp.dot(pn, ab)
@@ -1734,7 +1786,7 @@ def plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3):
# clip a polygon against another polygon # clip a polygon against another polygon
@wp.func @wp.func
def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n: wp.vec3, dir: wp.vec3): def _polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n: wp.vec3, dir: wp.vec3):
witness1 = mat3c() witness1 = mat3c()
witness2 = mat3c() witness2 = mat3c()
@@ -1746,10 +1798,10 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
pn = polyverts() pn = polyverts()
pd = polyvec() pd = polyvec()
for i in range(nface1 - 1): for i in range(nface1 - 1):
pdi, pni = plane_normal(face1[i], face1[i + 1], n) pdi, pni = _plane_normal(face1[i], face1[i + 1], n)
pd[i] = pdi pd[i] = pdi
pn[i] = pni pn[i] = pni
pdi, pni = plane_normal(face1[nface1 - 1], face1[0], n) pdi, pni = _plane_normal(face1[nface1 - 1], face1[0], n)
pd[nface1 - 1] = pdi pd[nface1 - 1] = pdi
pn[nface1 - 1] = pni pn[nface1 - 1] = pni
@@ -1770,8 +1822,8 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
Q = wp.where(i < npolygon - 1, polygon[i + 1], polygon[0]) Q = wp.where(i < npolygon - 1, polygon[i + 1], polygon[0])
# determine if P and Q are in the halfspace of the clipping edge # determine if P and Q are in the halfspace of the clipping edge
inside1 = halfspace(face1[e], pn[e], P) inside1 = _halfspace(face1[e], pn[e], P)
inside2 = halfspace(face1[e], pn[e], Q) inside2 = _halfspace(face1[e], pn[e], Q)
# PQ entirely outside the clipping edge, skip # PQ entirely outside the clipping edge, skip
if not inside1 and not inside2: if not inside1 and not inside2:
@@ -1784,7 +1836,7 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
continue continue
# add new vertex to clipped polygon where PQ intersects the clipping edge # add new vertex to clipped polygon where PQ intersects the clipping edge
t, res = plane_intersect(pn[e], pd[e], P, Q) t, res = _plane_intersect(pn[e], pd[e], P, Q)
if t >= 0.0 and t <= 1.0: if t >= 0.0 and t <= 1.0:
clipped[nclipped] = res clipped[nclipped] = res
nclipped += 1 nclipped += 1
@@ -1804,6 +1856,13 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
if npolygon < 1: if npolygon < 1:
return 0, witness1, witness2 return 0, witness1, witness2
if npolygon > 4:
quad = _polygon_quad(polygon, npolygon)
for i in range(4):
witness2[i] = polygon[quad[i]]
witness1[i] = witness2[i] - dir
return 4, witness1, witness2
# no pruning needed # no pruning needed
for i in range(npolygon): for i in range(npolygon):
witness2[i] = polygon[i] witness2[i] = polygon[i]
@@ -1813,7 +1872,7 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
# recover multiple contacts from EPA polytope # recover multiple contacts from EPA polytope
@wp.func @wp.func
def multicontact( def _multicontact(
pt: Polytope, face: wp.vec3i, x1: wp.vec3, x2: wp.vec3, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int pt: Polytope, face: wp.vec3i, x1: wp.vec3, x2: wp.vec3, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int
): ):
witness1 = mat3c() witness1 = mat3c()
@@ -1845,40 +1904,40 @@ def multicontact(
polymap = geom2.mesh_polymap polymap = geom2.mesh_polymap
# get dimensions of features of geoms 1 and 2 # get dimensions of features of geoms 1 and 2
nface1, feature_index1, feature_vertex1 = feature_dim(face, pt.vert_index1, pt.vert1) nface1, feature_index1, feature_vertex1 = _feature_dim(face, pt.vert_index1, pt.vert1)
nface2, feature_index2, feature_vertex2 = feature_dim(face, pt.vert_index2, pt.vert2) nface2, feature_index2, feature_vertex2 = _feature_dim(face, pt.vert_index2, pt.vert2)
dir = x2 - x1 dir = x2 - x1
dir_neg = -dir dir_neg = -dir
# get all possible face normals for each geom # get all possible face normals for each geom
if geomtype1 == int(GeomType.BOX.value): if geomtype1 == int(GeomType.BOX.value):
nnorms1, n1, idx1 = box_normals(nface1, feature_index1, geom1.rot, dir_neg) nnorms1, n1, idx1 = _box_normals(nface1, feature_index1, geom1.rot, dir_neg)
elif geomtype1 == int(GeomType.MESH.value): elif geomtype1 == int(GeomType.MESH.value):
nnorms1, n1, idx1 = mesh_normals( nnorms1, n1, idx1 = _mesh_normals(
nface1, feature_index1, geom1.rot, geom1.vertadr, geom1.mesh_polyadr, polynormal, polymapadr, polymapnum, polymap nface1, feature_index1, geom1.rot, geom1.vertadr, geom1.mesh_polyadr, polynormal, polymapadr, polymapnum, polymap
) )
if geomtype2 == int(GeomType.BOX.value): if geomtype2 == int(GeomType.BOX.value):
nnorms2, n2, idx2 = box_normals(nface2, feature_index2, geom2.rot, dir) nnorms2, n2, idx2 = _box_normals(nface2, feature_index2, geom2.rot, dir)
elif geomtype2 == int(GeomType.MESH.value): elif geomtype2 == int(GeomType.MESH.value):
nnorms2, n2, idx2 = mesh_normals( nnorms2, n2, idx2 = _mesh_normals(
nface2, feature_index2, geom2.rot, geom2.vertadr, geom2.mesh_polyadr, polynormal, polymapadr, polymapnum, polymap nface2, feature_index2, geom2.rot, geom2.vertadr, geom2.mesh_polyadr, polynormal, polymapadr, polymapnum, polymap
) )
# determine if any two face normals match # determine if any two face normals match
is_edge_contact_geom1 = 0 is_edge_contact_geom1 = 0
is_edge_contact_geom2 = 0 is_edge_contact_geom2 = 0
nres, res = aligned_faces(n1, nnorms1, n2, nnorms2) nres, res = _aligned_faces(n1, nnorms1, n2, nnorms2)
if not nres: if not nres:
# check if edge-face collision # check if edge-face collision
if nface1 < 3 and nface1 <= nface2: if nface1 < 3 and nface1 <= nface2:
nnorms1 = 0 nnorms1 = 0
if geomtype1 == int(GeomType.BOX.value): if geomtype1 == int(GeomType.BOX.value):
nnorms1, n1, endverts = box_edge_normals( nnorms1, n1, endverts = _box_edge_normals(
nface1, geom1.rot, geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0] nface1, geom1.rot, geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0]
) )
elif geomtype1 == int(GeomType.MESH.value): elif geomtype1 == int(GeomType.MESH.value):
nnorms1, n1, endverts = mesh_edge_normals( nnorms1, n1, endverts = _mesh_edge_normals(
nface1, nface1,
geom1.rot, geom1.rot,
geom1.pos, geom1.pos,
@@ -1895,7 +1954,7 @@ def multicontact(
feature_vertex1[1], feature_vertex1[1],
feature_index1[0], feature_index1[0],
) )
nres, res = aligned_face_edge(n1, nnorms1, n2, nnorms2) nres, res = _aligned_face_edge(n1, nnorms1, n2, nnorms2)
if not nres: if not nres:
return 1, witness1, witness2 return 1, witness1, witness2
is_edge_contact_geom1 = 1 is_edge_contact_geom1 = 1
@@ -1904,11 +1963,11 @@ def multicontact(
elif nface2 < 3: elif nface2 < 3:
nnorms2 = 0 nnorms2 = 0
if geomtype2 == int(GeomType.BOX.value): if geomtype2 == int(GeomType.BOX.value):
nnorms2, n2, endverts = box_edge_normals( nnorms2, n2, endverts = _box_edge_normals(
nface2, geom2.rot, geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0] nface2, geom2.rot, geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0]
) )
elif geomtype2 == int(GeomType.MESH.value): elif geomtype2 == int(GeomType.MESH.value):
nnorms2, n2, endverts = mesh_edge_normals( nnorms2, n2, endverts = _mesh_edge_normals(
nface2, nface2,
geom2.rot, geom2.rot,
geom2.pos, geom2.pos,
@@ -1925,7 +1984,7 @@ def multicontact(
feature_vertex2[1], feature_vertex2[1],
feature_index2[0], feature_index2[0],
) )
nres, res = aligned_face_edge(n2, nnorms2, n1, nnorms1) nres, res = _aligned_face_edge(n2, nnorms2, n1, nnorms1)
if not nres: if not nres:
return 1, witness1, witness2 return 1, witness1, witness2
is_edge_contact_geom2 = 1 is_edge_contact_geom2 = 1
@@ -1944,9 +2003,9 @@ def multicontact(
else: else:
ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i]) ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i])
if geomtype1 == int(GeomType.BOX.value): if geomtype1 == int(GeomType.BOX.value):
nface1, face1 = box_face(geom1.rot, geom1.pos, geom1.size, ind) nface1, face1 = _box_face(geom1.rot, geom1.pos, geom1.size, ind)
elif geomtype1 == int(GeomType.MESH.value): elif geomtype1 == int(GeomType.MESH.value):
nface1, face1 = mesh_face( nface1, face1 = _mesh_face(
geom1.rot, geom1.pos, geom1.vertadr, geom1.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, ind geom1.rot, geom1.pos, geom1.vertadr, geom1.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, ind
) )
@@ -1957,9 +2016,9 @@ def multicontact(
nface2 = 2 nface2 = 2
else: else:
if geomtype2 == int(GeomType.BOX.value): if geomtype2 == int(GeomType.BOX.value):
nface2, face2 = box_face(geom2.rot, geom2.pos, geom2.size, idx2[j]) nface2, face2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j])
elif geomtype2 == int(GeomType.MESH.value): elif geomtype2 == int(GeomType.MESH.value):
nface2, face2 = mesh_face( nface2, face2 = _mesh_face(
geom2.rot, geom2.pos, geom2.vertadr, geom2.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, idx2[j] geom2.rot, geom2.pos, geom2.vertadr, geom2.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, idx2[j]
) )
@@ -1971,20 +2030,20 @@ def multicontact(
# face1 is an edge; clip face1 against face2 # face1 is an edge; clip face1 against face2
if is_edge_contact_geom1: if is_edge_contact_geom1:
approx_dir = wp.norm_l2(dir) * n2[j] approx_dir = wp.norm_l2(dir) * n2[j]
return polygon_clip(face2, nface2, face1, nface1, n2[j], approx_dir) return _polygon_clip(face2, nface2, face1, nface1, n2[j], approx_dir)
# face2 is an edge; clip face2 against face1 # face2 is an edge; clip face2 against face1
if is_edge_contact_geom2: if is_edge_contact_geom2:
approx_dir = -wp.norm_l2(dir) * n1[j] approx_dir = -wp.norm_l2(dir) * n1[j]
return polygon_clip(face1, nface1, face2, nface2, n1[j], approx_dir) return _polygon_clip(face1, nface1, face2, nface2, n1[j], approx_dir)
# face-face collision # face-face collision
approx_dir = wp.norm_l2(dir) * n2[j] approx_dir = wp.norm_l2(dir) * n2[j]
return polygon_clip(face1, nface1, face2, nface2, n1[i], approx_dir) return _polygon_clip(face1, nface1, face2, nface2, n1[i], approx_dir)
@wp.func @wp.func
def inflate(dist: float, x1: wp.vec3, x2: wp.vec3, margin1: float, margin2: float): def _inflate(dist: float, x1: wp.vec3, x2: wp.vec3, margin1: float, margin2: float):
n = wp.normalize(x2 - x1) n = wp.normalize(x2 - x1)
if margin1 > 0.0: if margin1 > 0.0:
x1 += margin1 * n x1 += margin1 * n
@@ -2021,9 +2080,9 @@ def ccd(
face_map: wp.array(dtype=int), face_map: wp.array(dtype=int),
horizon: wp.array(dtype=int), horizon: wp.array(dtype=int),
): ):
"""General convex collision detection via GJK/EPA."""
witness1 = mat3c() witness1 = mat3c()
witness2 = mat3c() witness2 = mat3c()
"""General convex collision detection via GJK/EPA."""
margin1 = 0.0 margin1 = 0.0
margin2 = 0.0 margin2 = 0.0
@@ -2044,7 +2103,7 @@ def ccd(
witness1[0] = result.x1 witness1[0] = result.x1
witness2[0] = result.x2 witness2[0] = result.x2
return result.dist, 1, witness1, witness2 return result.dist, 1, witness1, witness2
dist, x1, x2 = inflate(result.dist, result.x1, result.x2, margin1, margin2) dist, x1, x2 = _inflate(result.dist, result.x1, result.x2, margin1, margin2)
witness1[0] = x1 witness1[0] = x1
witness2[0] = x2 witness2[0] = x2
return dist, 1, witness1, witness2 return dist, 1, witness1, witness2
@@ -2151,7 +2210,7 @@ def ccd(
and (geomtype1 == int(GeomType.BOX.value) or (geomtype1 == int(GeomType.MESH.value) and geom1.mesh_polyadr > -1)) and (geomtype1 == int(GeomType.BOX.value) or (geomtype1 == int(GeomType.MESH.value) and geom1.mesh_polyadr > -1))
and (geomtype2 == int(GeomType.BOX.value) or (geomtype2 == int(GeomType.MESH.value) and geom2.mesh_polyadr > -1)) and (geomtype2 == int(GeomType.BOX.value) or (geomtype2 == int(GeomType.MESH.value) and geom2.mesh_polyadr > -1))
): ):
num, w1, w2 = multicontact(pt, pt.face[idx], x1, x2, geom1, geom2, geomtype1, geomtype2) num, w1, w2 = _multicontact(pt, pt.face[idx], x1, x2, geom1, geom2, geomtype1, geomtype2)
if num > 0: if num > 0:
return dist, num, w1, w2 return dist, num, w1, w2
witness1[0] = x1 witness1[0] = x1
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py
Vendored
+2 -2
View File
@@ -424,7 +424,7 @@ class GJKTest(absltest.TestCase):
) )
_, ncon, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=True) _, ncon, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=True)
self.assertEqual(ncon, 5) self.assertEqual(ncon, 4)
def test_box_box_ccd2(self): def test_box_box_ccd2(self):
"""Test box-box multiccd 2.""" """Test box-box multiccd 2."""
@@ -441,7 +441,7 @@ class GJKTest(absltest.TestCase):
) )
_, ncon, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=True) _, ncon, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=True)
self.assertEqual(ncon, 5) self.assertEqual(ncon, 4)
if __name__ == "__main__": if __name__ == "__main__":
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+1 -18
View File
@@ -129,18 +129,6 @@ def _qderiv_tendon_damping(
qM_integration_out[worldid, dofjid, dofiid] -= qderiv qM_integration_out[worldid, dofjid, dofiid] -= qderiv
@wp.kernel
def _qfrc_forward(
# Data in:
qfrc_smooth_in: wp.array2d(dtype=float),
qfrc_constraint_in: wp.array2d(dtype=float),
# Data out:
qfrc_integration_out: wp.array2d(dtype=float),
):
worldid, dofid = wp.tid()
qfrc_integration_out[worldid, dofid] = qfrc_smooth_in[worldid, dofid] + qfrc_constraint_in[worldid, dofid]
@event_scope @event_scope
def deriv_smooth_vel(m: Model, d: Data, flg_forward: bool = True): def deriv_smooth_vel(m: Model, d: Data, flg_forward: bool = True):
"""Analytical derivative of smooth forces w.r.t. velocities. """Analytical derivative of smooth forces w.r.t. velocities.
@@ -193,12 +181,7 @@ def deriv_smooth_vel(m: Model, d: Data, flg_forward: bool = True):
) )
if flg_forward: if flg_forward:
wp.launch( wp.copy(d.qfrc_integration, d.efc.Ma)
_qfrc_forward,
dim=(d.nworld, m.nv),
inputs=[d.qfrc_smooth, d.qfrc_constraint],
outputs=[d.qfrc_integration],
)
else: else:
# qfrc = qM @ qacc # qfrc = qM @ qacc
mul_m(m, d, d.qfrc_integration, d.qacc, d.inverse_mul_m_skip, d.qM_integration) mul_m(m, d, d.qfrc_integration, d.qacc, d.inverse_mul_m_skip, d.qM_integration)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+6 -19
View File
@@ -319,11 +319,7 @@ def _euler_damp_qfrc_sparse(
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array(dtype=float),
dof_Madr: wp.array(dtype=int), dof_Madr: wp.array(dtype=int),
dof_damping: wp.array2d(dtype=float), dof_damping: wp.array2d(dtype=float),
# Data in:
qfrc_smooth_in: wp.array2d(dtype=float),
qfrc_constraint_in: wp.array2d(dtype=float),
# Data out: # Data out:
qfrc_integration_out: wp.array2d(dtype=float),
qM_integration_out: wp.array3d(dtype=float), qM_integration_out: wp.array3d(dtype=float),
): ):
worldid, tid = wp.tid() worldid, tid = wp.tid()
@@ -331,7 +327,6 @@ def _euler_damp_qfrc_sparse(
adr = dof_Madr[tid] adr = dof_Madr[tid]
qM_integration_out[worldid, 0, adr] += timestep * dof_damping[worldid, tid] qM_integration_out[worldid, 0, adr] += timestep * dof_damping[worldid, tid]
qfrc_integration_out[worldid, tid] = qfrc_smooth_in[worldid, tid] + qfrc_constraint_in[worldid, tid]
def _euler_sparse(m: Model, d: Data): def _euler_sparse(m: Model, d: Data):
@@ -343,11 +338,8 @@ def _euler_sparse(m: Model, d: Data):
m.opt.timestep, m.opt.timestep,
m.dof_Madr, m.dof_Madr,
m.dof_damping, m.dof_damping,
d.qfrc_smooth,
d.qfrc_constraint,
], ],
outputs=[ outputs=[
d.qfrc_integration,
d.qM_integration, d.qM_integration,
], ],
) )
@@ -358,7 +350,7 @@ def _euler_sparse(m: Model, d: Data):
d.qLD_integration, d.qLD_integration,
d.qLDiagInv_integration, d.qLDiagInv_integration,
d.qacc_integration, d.qacc_integration,
d.qfrc_integration, d.efc.Ma,
) )
@@ -371,8 +363,7 @@ def _tile_euler_dense(tile: TileSet):
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array(dtype=float),
# Data in: # Data in:
qM_in: wp.array3d(dtype=float), qM_in: wp.array3d(dtype=float),
qfrc_smooth_in: wp.array2d(dtype=float), efc_Ma_in: wp.array2d(dtype=float),
qfrc_constraint_in: wp.array2d(dtype=float),
# In: # In:
adr_in: wp.array(dtype=int), adr_in: wp.array(dtype=int),
# Data out: # Data out:
@@ -388,14 +379,10 @@ def _tile_euler_dense(tile: TileSet):
damping_scaled = damping_tile * timestep damping_scaled = damping_tile * timestep
qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled) qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled)
qfrc_smooth_tile = wp.tile_load(qfrc_smooth_in[worldid], shape=(TILE_SIZE,), offset=(dofid,)) Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
qfrc_constraint_tile = wp.tile_load(qfrc_constraint_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
qfrc_tile = qfrc_smooth_tile + qfrc_constraint_tile
L_tile = wp.tile_cholesky(qm_integration_tile) L_tile = wp.tile_cholesky(qm_integration_tile)
qacc_tile = wp.tile_cholesky_solve(L_tile, qfrc_tile) qacc_integration_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
wp.tile_store(qacc_integration_out[worldid], qacc_tile, offset=(dofid)) wp.tile_store(qacc_integration_out[worldid], qacc_integration_tile, offset=(dofid))
return euler_dense return euler_dense
@@ -413,7 +400,7 @@ def euler(m: Model, d: Data):
wp.launch_tiled( wp.launch_tiled(
_tile_euler_dense(tile), _tile_euler_dense(tile),
dim=(d.nworld, tile.adr.size), dim=(d.nworld, tile.adr.size),
inputs=[m.dof_damping, m.opt.timestep, d.qM, d.qfrc_smooth, d.qfrc_constraint, tile.adr], inputs=[m.dof_damping, m.opt.timestep, d.qM, d.efc.Ma, tile.adr],
outputs=[d.qacc_integration], outputs=[d.qacc_integration],
block_dim=m.block_dim.euler_dense, block_dim=m.block_dim.euler_dense,
) )
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py
Vendored
+10 -6
View File
@@ -25,7 +25,6 @@ import mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import test_util from mujoco.mjx.third_party.mujoco_warp._src import test_util
from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GainType from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
@@ -133,14 +132,13 @@ class ForwardTest(parameterized.TestCase):
mjm, mjd, _, _ = test_util.fixture("pendula.xml", kick=True, eulerdamp=eulerdamp, sparse=sparse) mjm, mjd, _, _ = test_util.fixture("pendula.xml", kick=True, eulerdamp=eulerdamp, sparse=sparse)
self.assertTrue((mjm.dof_damping > 0).any()) self.assertTrue((mjm.dof_damping > 0).any())
mjd.qvel[:] = 1.0
mjd.qacc[:] = 1.0
mujoco.mj_forward(mjm, mjd)
m = mjwarp.put_model(mjm) m = mjwarp.put_model(mjm)
d = mjwarp.put_data(mjm, mjd) d = mjwarp.put_data(mjm, mjd)
mujoco.mj_forward(mjm, mjd)
mujoco.mj_Euler(mjm, mjd) mujoco.mj_Euler(mjm, mjd)
mjwarp.solve(m, d) # compute efc.Ma
mjwarp.euler(m, d) mjwarp.euler(m, d)
_assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos") _assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
@@ -220,9 +218,11 @@ class ForwardTest(parameterized.TestCase):
m = mjwarp.put_model(mjm) m = mjwarp.put_model(mjm)
d = mjwarp.put_data(mjm, mjd) d = mjwarp.put_data(mjm, mjd)
mjwarp.implicit(m, d)
mujoco.mj_implicit(mjm, mjd) mujoco.mj_implicit(mjm, mjd)
mjwarp.solve(m, d) # compute efc.Ma
mjwarp.implicit(m, d)
_assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos") _assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
_assert_eq(d.act.numpy()[0], mjd.act, "act") _assert_eq(d.act.numpy()[0], mjd.act, "act")
@@ -230,6 +230,8 @@ class ForwardTest(parameterized.TestCase):
mjm, mjd, m, d = test_util.fixture("actuation/position.xml", keyframe=0, integrator=IntegratorType.IMPLICITFAST, kick=True) mjm, mjd, m, d = test_util.fixture("actuation/position.xml", keyframe=0, integrator=IntegratorType.IMPLICITFAST, kick=True)
mujoco.mj_implicit(mjm, mjd) mujoco.mj_implicit(mjm, mjd)
mjwarp.solve(m, d) # compute efc.Ma
mjwarp.implicit(m, d) mjwarp.implicit(m, d)
_assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos") _assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
@@ -239,6 +241,8 @@ class ForwardTest(parameterized.TestCase):
mjm, mjd, m, d = test_util.fixture("tendon/damping.xml", keyframe=0, integrator=IntegratorType.IMPLICITFAST, kick=True) mjm, mjd, m, d = test_util.fixture("tendon/damping.xml", keyframe=0, integrator=IntegratorType.IMPLICITFAST, kick=True)
mujoco.mj_implicit(mjm, mjd) mujoco.mj_implicit(mjm, mjd)
mjwarp.solve(m, d) # compute efc.Ma
mjwarp.implicit(m, d) mjwarp.implicit(m, d)
_assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos") _assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+3 -2
View File
@@ -129,8 +129,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
raise NotImplementedError("Contact sensor: only geom1-geom2 matching is implemented.") raise NotImplementedError("Contact sensor: only geom1-geom2 matching is implemented.")
# reduction # reduction
if (mjm.sensor_intprm[is_contact_sensor, 1] != 1).any(): if (~((mjm.sensor_intprm[is_contact_sensor, 1] == 1) | (mjm.sensor_intprm[is_contact_sensor, 1] == 2))).any():
raise NotImplementedError(f"Contact sensor: only mindist reduction is implemented.") raise NotImplementedError(f"Contact sensor: only mindist and maxforce reduction are implemented.")
# TODO(team): remove after _update_gradient for Newton uses tile operations for islands # TODO(team): remove after _update_gradient for Newton uses tile operations for islands
nv_max = 60 nv_max = 60
@@ -462,6 +462,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
impratio=create_nmodel_batched_array(np.array(mjm.opt.impratio), dtype=float, expand_dim=False), impratio=create_nmodel_batched_array(np.array(mjm.opt.impratio), dtype=float, expand_dim=False),
is_sparse=bool(is_sparse), is_sparse=bool(is_sparse),
ls_parallel=False, ls_parallel=False,
ls_parallel_min_step=1.0e-6, # TODO(team): determine good default setting
gjk_iterations=MJ_CCD_ITERATIONS, gjk_iterations=MJ_CCD_ITERATIONS,
epa_iterations=MJ_CCD_ITERATIONS, epa_iterations=MJ_CCD_ITERATIONS,
broadphase=int(broadphase), broadphase=int(broadphase),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io_test.py
Vendored
-1
View File
@@ -366,7 +366,6 @@ class IOTest(parameterized.TestCase):
'<contact geom1="plane"/>', '<contact geom1="plane"/>',
'<contact geom2="plane"/>', '<contact geom2="plane"/>',
'<contact site="site"/>', '<contact site="site"/>',
'<contact reduce="maxforce"/>',
'<contact reduce="netforce"/>', '<contact reduce="netforce"/>',
'<contact geom1="plane" geom2="sphere"/>', '<contact geom1="plane" geom2="sphere"/>',
) )
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+15 -3
View File
@@ -477,11 +477,23 @@ def _flex_bending(
if v[3] == -1: if v[3] == -1:
return return
frc = wp.mat(0.0, shape=(4, 3))
if flex_bending[17 * edgeid + 16]:
v0 = flexvert_xpos_in[worldid, v[0]]
v1 = flexvert_xpos_in[worldid, v[1]]
v2 = flexvert_xpos_in[worldid, v[2]]
v3 = flexvert_xpos_in[worldid, v[3]]
frc[1] = wp.cross(v2 - v0, v3 - v0)
frc[2] = wp.cross(v3 - v0, v1 - v0)
frc[3] = wp.cross(v1 - v0, v2 - v0)
frc[0] = -(frc[1] + frc[2] + frc[3])
force = wp.mat(0.0, shape=(nvert, 3)) force = wp.mat(0.0, shape=(nvert, 3))
for i in range(nvert): for i in range(nvert):
for j in range(nvert): for x in range(3):
for x in range(3): for j in range(nvert):
force[i, x] -= flex_bending[17*edgeid + 4*i + j] * flexvert_xpos_in[worldid, v[j]][x] force[i, x] -= flex_bending[17 * edgeid + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] -= flex_bending[17 * edgeid + 16] * frc[i, x]
for i in range(nvert): for i in range(nvert):
bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]] bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+6 -6
View File
@@ -510,20 +510,20 @@ def _ray_hfield(
y = safe_div(lpnt[1] + all[i] * lvec[1] + size[1], dy) y = safe_div(lpnt[1] + all[i] * lvec[1] + size[1], dy)
y0 = wp.max(0.0, wp.min(float(nrow - 2), wp.floor(y))) y0 = wp.max(0.0, wp.min(float(nrow - 2), wp.floor(y)))
if i == 1: if i == 1:
z0 = hfield_data[adr + int(wp.round(y0)) * nrow + ncol - 1] z0 = hfield_data[adr + int(wp.round(y0 + 0.0)) * ncol + ncol - 1]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * nrow + ncol - 1] z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * ncol + ncol - 1]
else: else:
z0 = hfield_data[adr + int(wp.round(y0)) * nrow] z0 = hfield_data[adr + int(wp.round(y0 + 0.0)) * ncol]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * nrow] z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * ncol]
# side normal to y-axis # side normal to y-axis
else: else:
y = safe_div(lpnt[0] + all[i] * lvec[0] + size[0], dx) y = safe_div(lpnt[0] + all[i] * lvec[0] + size[0], dx)
y0 = wp.max(0.0, wp.min(float(ncol - 2), wp.floor(y))) y0 = wp.max(0.0, wp.min(float(ncol - 2), wp.floor(y)))
if i == 3: if i == 3:
z0 = hfield_data[adr + int(wp.round(y0)) + (nrow - 1) * ncol] z0 = hfield_data[adr + int(wp.round(y0 + 0.0)) + (nrow - 1) * ncol]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) + (nrow - 1) * ncol] z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) + (nrow - 1) * ncol]
else: else:
z0 = hfield_data[adr + int(wp.round(y0))] z0 = hfield_data[adr + int(wp.round(y0 + 0.0))]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0))] z1 = hfield_data[adr + int(wp.round(y0 + 1.0))]
# check if point is below line segments # check if point is below line segments
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+40 -4
View File
@@ -1986,14 +1986,22 @@ def _sensor_tactile(
@wp.kernel @wp.kernel
def _contact_match( def _contact_match(
# Model: # Model:
opt_cone: int,
sensor_objid: wp.array(dtype=int), sensor_objid: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int), sensor_refid: wp.array(dtype=int),
sensor_intprm: wp.array2d(dtype=int),
sensor_contact_adr: wp.array(dtype=int), sensor_contact_adr: wp.array(dtype=int),
# Data in: # Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int), ncon_in: wp.array(dtype=int),
contact_dist_in: wp.array(dtype=float), contact_dist_in: wp.array(dtype=float),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int),
contact_geom_in: wp.array(dtype=wp.vec2i), contact_geom_in: wp.array(dtype=wp.vec2i),
contact_efc_address_in: wp.array2d(dtype=int),
contact_worldid_in: wp.array(dtype=int), contact_worldid_in: wp.array(dtype=int),
efc_force_in: wp.array2d(dtype=float),
# Data out: # Data out:
sensor_contact_nmatch_out: wp.array2d(dtype=int), sensor_contact_nmatch_out: wp.array2d(dtype=int),
sensor_contact_matchid_out: wp.array3d(dtype=int), sensor_contact_matchid_out: wp.array3d(dtype=int),
@@ -2009,6 +2017,7 @@ def _contact_match(
# sensor information # sensor information
objid = sensor_objid[sensorid] objid = sensor_objid[sensorid]
refid = sensor_refid[sensorid] refid = sensor_refid[sensorid]
reduce = sensor_intprm[sensorid, 1]
# contact information # contact information
geom = contact_geom_in[contactid] geom = contact_geom_in[contactid]
@@ -2022,8 +2031,27 @@ def _contact_match(
contactmatchid = wp.atomic_add(sensor_contact_nmatch_out[worldid], contactsensorid, 1) contactmatchid = wp.atomic_add(sensor_contact_nmatch_out[worldid], contactsensorid, 1)
sensor_contact_matchid_out[worldid, contactsensorid, contactmatchid] = contactid sensor_contact_matchid_out[worldid, contactsensorid, contactmatchid] = contactid
# TODO(thowell): alternative criteria if reduce == 1: # mindist
sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = contact_dist_in[contactid] sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = contact_dist_in[contactid]
elif reduce == 2: # maxforce
contact_force = support.contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
contact_efc_address_in,
efc_force_in,
worldid,
contactid,
False,
)
force_magnitude = (
contact_force[0] * contact_force[0] + contact_force[1] * contact_force[1] + contact_force[2] * contact_force[2]
)
sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = -force_magnitude
# TODO(thowell): netforce
# contact direction # contact direction
if geom1geom0: if geom1geom0:
@@ -2155,20 +2183,28 @@ def sensor_acc(m: Model, d: Data):
if m.sensor_contact_adr.size: if m.sensor_contact_adr.size:
# match criteria # match criteria
d.sensor_contact_nmatch.zero_() d.sensor_contact_nmatch.zero_()
d.sensor_contact_matchid.zero_() d.sensor_contact_matchid.fill_(-1)
d.sensor_contact_criteria.zero_() d.sensor_contact_criteria.fill_(1.0e32)
wp.launch( wp.launch(
_contact_match, _contact_match,
dim=(m.sensor_contact_adr.size, d.nconmax), dim=(m.sensor_contact_adr.size, d.nconmax),
inputs=[ inputs=[
m.opt.cone,
m.sensor_objid, m.sensor_objid,
m.sensor_refid, m.sensor_refid,
m.sensor_intprm,
m.sensor_contact_adr, m.sensor_contact_adr,
d.njmax,
d.ncon, d.ncon,
d.contact.dist, d.contact.dist,
d.contact.frame,
d.contact.friction,
d.contact.dim,
d.contact.geom, d.contact.geom,
d.contact.efc_address,
d.contact.worldid, d.contact.worldid,
d.efc.force,
], ],
outputs=[ outputs=[
d.sensor_contact_nmatch, d.sensor_contact_nmatch,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor_test.py
Vendored
+12 -16
View File
@@ -411,7 +411,7 @@ class SensorTest(parameterized.TestCase):
@parameterized.parameters( @parameterized.parameters(
'type="sphere" size=".1"', 'type="sphere" size=".1"',
'type="capsule" size=".1 .1" euler="0 89 89"', 'type="capsule" size=".1 .1" euler="0 89 89"',
'type="box" size=".1 .1 .1" euler=".02 .05 .1"', 'type="box" size=".1 .11 .12" euler=".02 .05 .1"',
) )
def test_contact_sensor(self, geom): def test_contact_sensor(self, geom):
"""Test contact sensor.""" """Test contact sensor."""
@@ -419,22 +419,18 @@ class SensorTest(parameterized.TestCase):
contact_sensor = "" contact_sensor = ""
# data combinations # data combinations
field = ["found", "force", "torque", "dist", "pos", "normal", "tangent"] datas = ["found", "force dist normal", "torque pos tangent", "found force torque dist pos normal tangent"]
datas = itertools.chain.from_iterable([itertools.combinations(field, i) for i in range(len(field))])
datas = list(datas)
for num in [1, 2, 3, 4, 5]: for geoms in [
for geoms in [ 'geom1="plane" geom2="geom"',
'geom1="plane" geom2="geom"', 'geom1="geom" geom2="plane"',
'geom1="geom" geom2="plane"', 'geom1="sphere" geom2="plane"',
'geom1="plane" geom2="sphere"', 'geom1="geom" geom2="sphere"',
'geom1="sphere" geom2="plane"', ]:
'geom1="geom" geom2="sphere"', for num in [1, 3, 5]:
'geom1="sphere" geom2="geom"', for reduce in ["mindist", "maxforce"]:
]: for data in datas:
for data in datas: contact_sensor += f'<contact {geoms} num="{num}" reduce="{reduce}" data="{data}"/>'
data = " ".join(data)
contact_sensor += f'<contact {geoms} num="{num}" reduce="mindist" data="{data}"/>'
_MJCF = f""" _MJCF = f"""
<mujoco> <mujoco>
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+92 -71
View File
@@ -720,11 +720,18 @@ def _linesearch_iterative(m: types.Model, d: types.Data):
) )
@wp.func
def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> float:
step = (wp.log(max_value) - wp.log(min_value)) / wp.max(1.0, float(num_values - 1))
return wp.exp(wp.log(min_value) + float(i) * step)
@wp.kernel @wp.kernel
def linesearch_parallel_fused( def linesearch_parallel_fused(
# Model: # Model:
nlsp: int, nlsp: int,
opt_impratio: wp.array(dtype=float), opt_impratio: wp.array(dtype=float),
opt_ls_parallel_min_step: float,
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), ncon_in: wp.array(dtype=int),
@@ -750,7 +757,7 @@ def linesearch_parallel_fused(
if efc_done_in[worldid]: if efc_done_in[worldid]:
return return
alpha = float(alphaid) / float(nlsp - 1) alpha = _log_scale(opt_ls_parallel_min_step, 1.0, nlsp, alphaid)
out = _eval_cost(efc_quad_gauss_in[worldid], alpha) out = _eval_cost(efc_quad_gauss_in[worldid], alpha)
@@ -846,6 +853,7 @@ def linesearch_parallel_fused(
def linesearch_parallel_best_alpha( def linesearch_parallel_best_alpha(
# Model: # Model:
nlsp: int, nlsp: int,
opt_ls_parallel_min_step: float,
# Data in: # Data in:
efc_done_in: wp.array(dtype=bool), efc_done_in: wp.array(dtype=bool),
efc_cost_candidate_in: wp.array2d(dtype=float), efc_cost_candidate_in: wp.array2d(dtype=float),
@@ -867,7 +875,7 @@ def linesearch_parallel_best_alpha(
best_cost = cost best_cost = cost
bestid = i bestid = i
efc_alpha_out[worldid] = float(bestid) / float(nlsp - 1) efc_alpha_out[worldid] = _log_scale(opt_ls_parallel_min_step, 1.0, nlsp, bestid)
def _linesearch_parallel(m: types.Model, d: types.Data): def _linesearch_parallel(m: types.Model, d: types.Data):
@@ -877,6 +885,7 @@ def _linesearch_parallel(m: types.Model, d: types.Data):
inputs=[ inputs=[
m.nlsp, m.nlsp,
m.opt.impratio, m.opt.impratio,
m.opt.ls_parallel_min_step,
d.njmax, d.njmax,
d.ncon, d.ncon,
d.ne, d.ne,
@@ -900,7 +909,7 @@ def _linesearch_parallel(m: types.Model, d: types.Data):
wp.launch( wp.launch(
linesearch_parallel_best_alpha, linesearch_parallel_best_alpha,
dim=(d.nworld), dim=(d.nworld),
inputs=[m.nlsp, d.efc.done, d.efc.cost_candidate], inputs=[m.nlsp, m.opt.ls_parallel_min_step, d.efc.done, d.efc.cost_candidate],
outputs=[d.efc.alpha], outputs=[d.efc.alpha],
) )
@@ -1589,26 +1598,6 @@ def update_gradient_grad(
wp.atomic_add(efc_grad_dot_out, worldid, grad * grad) wp.atomic_add(efc_grad_dot_out, worldid, grad * grad)
@wp.kernel
def update_gradient_zero_h_lower(
# Model:
dof_tri_row: wp.array(dtype=int),
dof_tri_col: wp.array(dtype=int),
# Data in:
efc_done_in: wp.array(dtype=bool),
# Data out:
efc_h_out: wp.array3d(dtype=float),
):
worldid, elementid = wp.tid()
if efc_done_in[worldid]:
return
rowid = dof_tri_row[elementid]
colid = dof_tri_col[elementid]
efc_h_out[worldid, rowid, colid] = 0.0
@wp.kernel @wp.kernel
def update_gradient_set_h_qM_lower_sparse( def update_gradient_set_h_qM_lower_sparse(
# Model: # Model:
@@ -1627,32 +1616,11 @@ def update_gradient_set_h_qM_lower_sparse(
i = qM_fullm_i[elementid] i = qM_fullm_i[elementid]
j = qM_fullm_j[elementid] j = qM_fullm_j[elementid]
efc_h_out[worldid, i, j] = qM_in[worldid, 0, elementid] efc_h_out[worldid, i, j] += qM_in[worldid, 0, elementid]
@wp.kernel @wp.kernel
def update_gradient_copy_lower_triangle( def update_gradient_JTDAJ_sparse(
# Model:
dof_tri_row: wp.array(dtype=int),
dof_tri_col: wp.array(dtype=int),
# Data in:
qM_in: wp.array3d(dtype=float),
efc_done_in: wp.array(dtype=bool),
# Data out:
efc_h_out: wp.array3d(dtype=float),
):
worldid, elementid = wp.tid()
if efc_done_in[worldid]:
return
rowid = dof_tri_row[elementid]
colid = dof_tri_col[elementid]
efc_h_out[worldid, rowid, colid] = qM_in[worldid, rowid, colid]
@wp.kernel
def update_gradient_JTDAJ(
# Data in: # Data in:
njmax_in: int, njmax_in: int,
nefc_in: wp.array(dtype=int), nefc_in: wp.array(dtype=int),
@@ -1673,12 +1641,14 @@ def update_gradient_JTDAJ(
dofi = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2 dofi = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
dofj = elementid - (dofi * (dofi + 1)) // 2 dofj = elementid - (dofi * (dofi + 1)) // 2
sum_h = float(0.0) # To optimize the loop, data for the next iteration is prefetched
# This allows to parallelize memory load and computation to hide memory latency
efc_state = efc_state_in[worldid, 0]
efc_D = efc_D_in[worldid, 0] efc_D = efc_D_in[worldid, 0]
# TODO(team): sparse efc_J # TODO(team): sparse efc_J
efc_Ji = efc_J_in[worldid, 0, dofi] efc_Ji = efc_J_in[worldid, 0, dofi]
efc_Jj = efc_J_in[worldid, 0, dofj] efc_Jj = efc_J_in[worldid, 0, dofj]
efc_state = efc_state_in[worldid, 0] sum_h = float(0.0)
for efcid in range(min(njmax_in, nefc) - 1): for efcid in range(min(njmax_in, nefc) - 1):
if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0: if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0:
sum_h += efc_Ji * efc_Jj * efc_D sum_h += efc_Ji * efc_Jj * efc_D
@@ -1689,9 +1659,60 @@ def update_gradient_JTDAJ(
efc_Jj = efc_J_in[worldid, jj, dofj] efc_Jj = efc_J_in[worldid, jj, dofj]
efc_state = efc_state_in[worldid, jj] efc_state = efc_state_in[worldid, jj]
# Adding the contribution from the last constraint row
if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0: if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0:
sum_h += efc_Ji * efc_Jj * efc_D sum_h += efc_Ji * efc_Jj * efc_D
efc_h_out[worldid, dofi, dofj] += sum_h
efc_h_out[worldid, dofi, dofj] = sum_h
@wp.kernel
def update_gradient_JTDAJ_dense(
# Data in:
njmax_in: int,
nefc_in: wp.array(dtype=int),
qM_in: wp.array3d(dtype=float),
efc_J_in: wp.array3d(dtype=float),
efc_D_in: wp.array2d(dtype=float),
efc_state_in: wp.array2d(dtype=int),
efc_done_in: wp.array(dtype=bool),
# Data out:
efc_h_out: wp.array3d(dtype=float),
):
worldid, elementid = wp.tid()
if efc_done_in[worldid]:
return
nefc = nefc_in[worldid]
dofi = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
dofj = elementid - (dofi * (dofi + 1)) // 2
# To optimize the loop, data for the next iteration is prefetched
# This allows to parallelize memory load and computation to hide memory latency
efc_state = efc_state_in[worldid, 0]
efc_D = efc_D_in[worldid, 0]
# TODO(team): sparse efc_J
efc_Ji = efc_J_in[worldid, 0, dofi]
efc_Jj = efc_J_in[worldid, 0, dofj]
sum_h = float(0.0)
for efcid in range(min(njmax_in, nefc) - 1):
if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0:
sum_h += efc_Ji * efc_Jj * efc_D
jj = efcid + 1
efc_D = efc_D_in[worldid, jj]
efc_Ji = efc_J_in[worldid, jj, dofi]
efc_Jj = efc_J_in[worldid, jj, dofj]
efc_state = efc_state_in[worldid, jj]
# Adding the contribution from the last constraint row
if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0:
sum_h += efc_Ji * efc_Jj * efc_D
qM = qM_in[worldid, dofi, dofj]
efc_h_out[worldid, dofi, dofj] = qM + sum_h
# TODO(thowell): combine with JTDAJ ? # TODO(thowell): combine with JTDAJ ?
@@ -1904,11 +1925,19 @@ def _update_gradient(m: types.Model, d: types.Data):
smooth.solve_m(m, d, d.efc.Mgrad, d.efc.grad) smooth.solve_m(m, d, d.efc.Mgrad, d.efc.grad)
elif m.opt.solver == types.SolverType.NEWTON: elif m.opt.solver == types.SolverType.NEWTON:
# h = qM + (efc_J.T * efc_D * active) @ efc_J # h = qM + (efc_J.T * efc_D * active) @ efc_J
lower_triangle_dim = int(m.nv * (m.nv + 1) / 2)
if m.opt.is_sparse: if m.opt.is_sparse:
wp.launch( wp.launch(
update_gradient_zero_h_lower, update_gradient_JTDAJ_sparse,
dim=(d.nworld, m.dof_tri_row.size), dim=(d.nworld, lower_triangle_dim),
inputs=[m.dof_tri_row, m.dof_tri_col, d.efc.done], inputs=[
d.njmax,
d.nefc,
d.efc.J,
d.efc.D,
d.efc.state,
d.efc.done,
],
outputs=[d.efc.h], outputs=[d.efc.h],
) )
wp.launch( wp.launch(
@@ -1919,28 +1948,20 @@ def _update_gradient(m: types.Model, d: types.Data):
) )
else: else:
wp.launch( wp.launch(
update_gradient_copy_lower_triangle, update_gradient_JTDAJ_dense,
dim=(d.nworld, m.dof_tri_row.size), dim=(d.nworld, lower_triangle_dim),
inputs=[m.dof_tri_row, m.dof_tri_col, d.qM, d.efc.done], inputs=[
d.njmax,
d.nefc,
d.qM,
d.efc.J,
d.efc.D,
d.efc.state,
d.efc.done,
],
outputs=[d.efc.h], outputs=[d.efc.h],
) )
lower_triangle_dim = int(m.nv * (m.nv + 1) / 2)
# TODO(team): Investigate whether d.efc.h initialization can be merged into this kernel
wp.launch(
update_gradient_JTDAJ,
dim=(d.nworld, lower_triangle_dim),
inputs=[
d.njmax,
d.nefc,
d.efc.J,
d.efc.D,
d.efc.state,
d.efc.done,
],
outputs=[d.efc.h],
)
if m.opt.cone == types.ConeType.ELLIPTIC: if m.opt.cone == types.ConeType.ELLIPTIC:
# Optimization: launching update_gradient_JTCJ with limited number of blocks on a GPU. # Optimization: launching update_gradient_JTCJ with limited number of blocks on a GPU.
# Profiling suggests that only a fraction of blocks out of the original # Profiling suggests that only a fraction of blocks out of the original
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver_test.py
Vendored
+28 -9
View File
@@ -42,8 +42,8 @@ def _assert_eq(a, b, name):
class SolverTest(parameterized.TestCase): class SolverTest(parameterized.TestCase):
@parameterized.product(cone=tuple(ConeType), solver_=tuple(SolverType)) @parameterized.product(cone=tuple(ConeType), solver_=tuple(SolverType))
def test_cost(self, cone, solver_): def test_constraint_update(self, cone, solver_):
"""Tests cost function is correct.""" """Tests _update_constraint function is correct."""
for keyframe in range(3): for keyframe in range(3):
mjm, mjd, m, d = test_util.fixture( mjm, mjd, m, d = test_util.fixture(
"constraints.xml", "constraints.xml",
@@ -60,14 +60,37 @@ class SolverTest(parameterized.TestCase):
mujoco.mj_constraintUpdate(mjm, mjd, jaref - mjd.efc_aref, cost, 0) mujoco.mj_constraintUpdate(mjm, mjd, jaref - mjd.efc_aref, cost, 0)
return cost return cost
mj_cost = cost(mjd.qacc) mjd_cost = cost(mjd.qacc)
# solve with 0 iterations just initializes constraints and costs and then exits # solve with 0 iterations just initializes constraints and costs and then exits
d.efc.force.zero_()
d.qfrc_constraint.zero_()
mjwarp.solve(m, d) mjwarp.solve(m, d)
mjwarp_cost = d.efc.cost.numpy()[0] - d.efc.gauss.numpy()[0] # Get the ordering indices based on efc_force, efc_state for MJWarp
nefc = d.nefc.numpy()[0]
efc_force = d.efc.force.numpy()[0, :nefc]
efc_state = d.efc.state.numpy()[0, :nefc]
# Get the ordering indices based on efc_force, efc_state for MuJoCo
mjd_efc_force = mjd.efc_force[:nefc]
mjd_efc_state = mjd.efc_state[:nefc]
_assert_eq(mjwarp_cost, mj_cost, name="cost") # Create sorting keys using lexsort (more efficient for multiple keys)
d_sort_indices = np.lexsort((efc_force, efc_state))
mjd_sort_indices = np.lexsort((mjd_efc_force, mjd_efc_state))
efc_cost = d.efc.cost.numpy()[0] - d.efc.gauss.numpy()[0]
qfrc_constraint = d.qfrc_constraint.numpy()[0]
efc_sorted_force = efc_force[d_sort_indices]
efc_sorted_state = efc_state[d_sort_indices]
mjd_sorted_force = mjd_efc_force[mjd_sort_indices]
mjd_sorted_state = mjd_efc_state[mjd_sort_indices]
_assert_eq(efc_sorted_state, mjd_sorted_state, "efc_state")
_assert_eq(efc_sorted_force, mjd_sorted_force, "efc_force")
_assert_eq(efc_cost, mjd_cost, "cost")
_assert_eq(qfrc_constraint, mjd.qfrc_constraint, "qfrc_constraint")
def test_init_linesearch(self): def test_init_linesearch(self):
"""Test linesearch initialization.""" """Test linesearch initialization."""
@@ -160,10 +183,6 @@ class SolverTest(parameterized.TestCase):
def test_parallel_linesearch(self, cone): def test_parallel_linesearch(self, cone):
"""Test that iterative and parallel linesearch leads to equivalent results.""" """Test that iterative and parallel linesearch leads to equivalent results."""
# TODO(team): Enable this case when elliptic/parallel linesearch is working
if cone == ConeType.ELLIPTIC:
return
_, _, m, d = test_util.fixture( _, _, m, d = test_util.fixture(
"humanoid/humanoid.xml", "humanoid/humanoid.xml",
cone=cone, cone=cone,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+3 -40
View File
@@ -167,45 +167,6 @@ def mul_m(
) )
@wp.kernel
def xfrc_accumulate_kernel(
# Model:
nbody: int,
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
dof_bodyid: wp.array(dtype=int),
# Data in:
xfrc_applied_in: wp.array2d(dtype=wp.spatial_vector),
xipos_in: wp.array2d(dtype=wp.vec3),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector),
# Out:
out: wp.array2d(dtype=float),
):
"""Accumulate applied forces on the subtree of a dof."""
worldid, dofid = wp.tid()
cdof = cdof_in[worldid, dofid]
rotational_cdof = wp.spatial_top(cdof)
jac = wp.spatial_vector(cdof[3], cdof[4], cdof[5], cdof[0], cdof[1], cdof[2])
bodyid = dof_bodyid[dofid]
accumul = float(0.0)
for child in range(bodyid, nbody):
# any body that is in the subtree of dof_bodyid is part of the jacobian
parentid = child
while parentid != 0 and parentid != bodyid:
parentid = body_parentid[parentid]
if parentid == 0:
continue # body is not part of the subtree
offset = xipos_in[worldid, child] - subtree_com_in[worldid, body_rootid[child]]
cross_term = wp.cross(rotational_cdof, offset)
xfrc_applied = xfrc_applied_in[worldid, child]
accumul += wp.dot(jac, xfrc_applied) + wp.dot(cross_term, wp.spatial_top(xfrc_applied))
out[worldid, dofid] += accumul
@wp.kernel @wp.kernel
def _apply_ft( def _apply_ft(
# Model: # Model:
@@ -232,6 +193,9 @@ def _apply_ft(
accumul = float(0.0) accumul = float(0.0)
for bodyid in range(dofbodyid, nbody): for bodyid in range(dofbodyid, nbody):
ft_body = ft_in[worldid, bodyid]
if ft_body == wp.spatial_vector():
continue
# any body that is in the subtree of dofbodyid is part of the jacobian # any body that is in the subtree of dofbodyid is part of the jacobian
parentid = bodyid parentid = bodyid
while parentid != 0 and parentid != dofbodyid: while parentid != 0 and parentid != dofbodyid:
@@ -240,7 +204,6 @@ def _apply_ft(
continue # body is not part of the subtree continue # body is not part of the subtree
offset = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]] offset = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]]
cross_term = wp.cross(rotational_cdof, offset) cross_term = wp.cross(rotational_cdof, offset)
ft_body = ft_in[worldid, bodyid]
accumul += wp.dot(jac, ft_body) + wp.dot(cross_term, wp.spatial_top(ft_body)) accumul += wp.dot(jac, ft_body) + wp.dot(cross_term, wp.spatial_top(ft_body))
if flg_add: if flg_add:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py
Vendored
+15 -13
View File
@@ -196,7 +196,7 @@ def benchmark(
event_trace: bool = False, event_trace: bool = False,
measure_alloc: bool = False, measure_alloc: bool = False,
measure_solver_niter: bool = False, measure_solver_niter: bool = False,
) -> Tuple[float, float, dict, list, list, list]: ) -> Tuple[float, float, dict, list, list, list, int]:
"""Benchmark a function of Model and Data. """Benchmark a function of Model and Data.
Args: Args:
@@ -217,22 +217,20 @@ def benchmark(
list: Number of contacts. list: Number of contacts.
list: Number of constraints. list: Number of constraints.
list: Number of solver iterations. list: Number of solver iterations.
int: Number of converged worlds.
""" """
jit_beg = time.perf_counter()
fn(m, d)
jit_end = time.perf_counter()
jit_duration = jit_end - jit_beg
wp.synchronize()
trace = {} trace = {}
ncon, nefc, solver_niter = [], [], [] ncon, nefc, solver_niter = [], [], []
with warp_util.EventTracer(enabled=event_trace) as tracer: with warp_util.EventTracer(enabled=event_trace) as tracer:
# capture the whole function as a CUDA graph # capture the whole function as a CUDA graph
jit_beg = time.perf_counter()
with wp.ScopedCapture() as capture: with wp.ScopedCapture() as capture:
fn(m, d) fn(m, d)
jit_end = time.perf_counter()
jit_duration = jit_end - jit_beg
graph = capture.graph graph = capture.graph
time_vec = np.zeros(nstep) time_vec = np.zeros(nstep)
@@ -258,14 +256,15 @@ def benchmark(
else: else:
trace = tracer.trace() trace = tracer.trace()
if measure_alloc: if measure_alloc:
ncon.append(d.ncon.numpy()[0]) ncon.append(np.max([d.ncon.numpy()[0], d.ncollision.numpy()[0]]))
nefc.append(np.sum(d.nefc.numpy())) nefc.append(np.max(d.nefc.numpy()))
if measure_solver_niter: if measure_solver_niter:
solver_niter.append(d.solver_niter.numpy()) solver_niter.append(d.solver_niter.numpy())
nsuccess = np.sum(~np.any(np.isnan(d.qpos.numpy()), axis=1))
run_duration = np.sum(time_vec) run_duration = np.sum(time_vec)
return jit_duration, run_duration, trace, ncon, nefc, solver_niter return jit_duration, run_duration, trace, ncon, nefc, solver_niter, nsuccess
class BenchmarkSuite: class BenchmarkSuite:
@@ -322,17 +321,20 @@ class BenchmarkSuite:
mujoco.mj_forward(mjm, mjd) mujoco.mj_forward(mjm, mjd)
wp.init() wp.init()
if os.environ.get("ASV_CACHE_KERNELS", "false").lower() == "false":
wp.clear_kernel_cache()
free_before = wp.get_device().free_memory free_before = wp.get_device().free_memory
m = io.put_model(mjm) m = io.put_model(mjm)
d = io.put_data(mjm, mjd, self.batch_size, self.nconmax, self.njmax) d = io.put_data(mjm, mjd, self.batch_size, self.nconmax, self.njmax)
free_after = wp.get_device().free_memory
jit_duration, _, trace, _, _, solver_niter = benchmark(forward.step, m, d, 1000, True, False, True) jit_duration, _, trace, _, _, solver_niter, _ = benchmark(forward.step, m, d, 1000, True, False, True)
metrics = { metrics = {
"jit_duration": jit_duration, "jit_duration": jit_duration,
"solver_niter_mean": np.mean(solver_niter), "solver_niter_mean": np.mean(solver_niter),
"solver_niter_p95": np.quantile(solver_niter, 0.95), "solver_niter_p95": np.quantile(solver_niter, 0.95),
"device_memory_allocated": free_before - wp.get_device().free_memory, "device_memory_allocated": free_before - free_after,
} }
def tree_flatten(d, parent_k=""): def tree_flatten(d, parent_k=""):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+3 -1
View File
@@ -557,6 +557,7 @@ class Option:
gjk_iterations: number of Gjk iterations in the convex narrowphase gjk_iterations: number of Gjk iterations in the convex narrowphase
epa_iterations: number of Epa iterations in the convex narrowphase epa_iterations: number of Epa iterations in the convex narrowphase
ls_parallel: evaluate engine solver step sizes in parallel ls_parallel: evaluate engine solver step sizes in parallel
ls_parallel_min_step: minimum step size for solver linesearch
wind: wind (for lift, drag, and viscosity) wind: wind (for lift, drag, and viscosity)
has_fluid: True if wind, density, or viscosity are non-zero at put_model time has_fluid: True if wind, density, or viscosity are non-zero at put_model time
density: density of medium density: density of medium
@@ -587,7 +588,8 @@ class Option:
is_sparse: bool is_sparse: bool
gjk_iterations: int # warp only gjk_iterations: int # warp only
epa_iterations: int # warp only epa_iterations: int # warp only
ls_parallel: bool ls_parallel: bool # warp only
ls_parallel_min_step: float # warp only
wind: wp.array(dtype=wp.vec3) wind: wp.array(dtype=wp.vec3)
has_fluid: bool has_fluid: bool
density: wp.array(dtype=float) density: wp.array(dtype=float)
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+6 -4
View File
@@ -50,6 +50,7 @@ _ENGINE = flags.DEFINE_enum_class("engine", EngineOptions.WARP, EngineOptions, "
_NCONMAX = flags.DEFINE_integer("nconmax", None, "Maximum number of contacts.") _NCONMAX = flags.DEFINE_integer("nconmax", None, "Maximum number of contacts.")
_NJMAX = flags.DEFINE_integer("njmax", None, "Maximum number of constraints per world.") _NJMAX = flags.DEFINE_integer("njmax", None, "Maximum number of constraints per world.")
_OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation: foo.bar = baz)", short_name="o") _OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation: foo.bar = baz)", short_name="o")
_KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
_DEVICE = flags.DEFINE_string("device", None, "override the default Warp device") _DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
@@ -100,7 +101,7 @@ def _override(model: Union[mjw.Model, mujoco.MjModel]):
if key in enum_fields: if key in enum_fields:
try: try:
val = str(enum_fields[key][val.upper()]) val = int(enum_fields[key][val.upper()])
except KeyError: except KeyError:
raise app.UsageError(f"Unrecognized enum value: {val}") raise app.UsageError(f"Unrecognized enum value: {val}")
@@ -110,13 +111,12 @@ def _override(model: Union[mjw.Model, mujoco.MjModel]):
raise app.UsageError(f"Unrecognized model field: {key}") raise app.UsageError(f"Unrecognized model field: {key}")
if i < len(attrs) - 1: if i < len(attrs) - 1:
obj = getattr(obj, attr) obj = getattr(obj, attr)
else: elif key not in enum_fields:
try: try:
val = type(getattr(obj, attr))(ast.literal_eval(val)) val = type(getattr(obj, attr))(ast.literal_eval(val))
except (SyntaxError, ValueError): except (SyntaxError, ValueError):
raise app.UsageError(f"Unrecognized value for field: {key}") raise app.UsageError(f"Unrecognized value for field: {key}")
setattr(obj, attr, val)
setattr(obj, attr, val)
def _compile_step(m, d): def _compile_step(m, d):
@@ -138,6 +138,8 @@ def _main(argv: Sequence[str]) -> None:
mjm = _load_model(epath.Path(argv[1])) mjm = _load_model(epath.Path(argv[1]))
mjd = mujoco.MjData(mjm) mjd = mujoco.MjData(mjm)
if mjm.nkey > 0 and _KEYFRAME.value > -1:
mujoco.mj_resetDataKeyframe(mjm, mjd, _KEYFRAME.value)
mujoco.mj_forward(mjm, mjd) mujoco.mj_forward(mjm, mjd)
if _ENGINE.value == EngineOptions.C: if _ENGINE.value == EngineOptions.C:
mjx/mujoco/mjx/warp/forward.py
+6
View File
@@ -333,6 +333,7 @@ def _forward_shim(
opt__iterations: int, opt__iterations: int,
opt__ls_iterations: int, opt__ls_iterations: int,
opt__ls_parallel: bool, opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
opt__ls_tolerance: wp.array(dtype=float), opt__ls_tolerance: wp.array(dtype=float),
opt__magnetic: wp.array(dtype=wp.vec3), opt__magnetic: wp.array(dtype=wp.vec3),
opt__run_collision_detection: bool, opt__run_collision_detection: bool,
@@ -734,6 +735,7 @@ def _forward_shim(
_m.opt.iterations = opt__iterations _m.opt.iterations = opt__iterations
_m.opt.ls_iterations = opt__ls_iterations _m.opt.ls_iterations = opt__ls_iterations
_m.opt.ls_parallel = opt__ls_parallel _m.opt.ls_parallel = opt__ls_parallel
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
_m.opt.ls_tolerance = opt__ls_tolerance _m.opt.ls_tolerance = opt__ls_tolerance
_m.opt.magnetic = opt__magnetic _m.opt.magnetic = opt__magnetic
_m.opt.run_collision_detection = opt__run_collision_detection _m.opt.run_collision_detection = opt__run_collision_detection
@@ -1662,6 +1664,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.opt.iterations, m.opt.iterations,
m.opt.ls_iterations, m.opt.ls_iterations,
m.opt._impl.ls_parallel, m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
m.opt.ls_tolerance, m.opt.ls_tolerance,
m.opt.magnetic, m.opt.magnetic,
m.opt._impl.run_collision_detection, m.opt._impl.run_collision_detection,
@@ -2358,6 +2361,7 @@ def _step_shim(
opt__iterations: int, opt__iterations: int,
opt__ls_iterations: int, opt__ls_iterations: int,
opt__ls_parallel: bool, opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
opt__ls_tolerance: wp.array(dtype=float), opt__ls_tolerance: wp.array(dtype=float),
opt__magnetic: wp.array(dtype=wp.vec3), opt__magnetic: wp.array(dtype=wp.vec3),
opt__run_collision_detection: bool, opt__run_collision_detection: bool,
@@ -2773,6 +2777,7 @@ def _step_shim(
_m.opt.iterations = opt__iterations _m.opt.iterations = opt__iterations
_m.opt.ls_iterations = opt__ls_iterations _m.opt.ls_iterations = opt__ls_iterations
_m.opt.ls_parallel = opt__ls_parallel _m.opt.ls_parallel = opt__ls_parallel
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
_m.opt.ls_tolerance = opt__ls_tolerance _m.opt.ls_tolerance = opt__ls_tolerance
_m.opt.magnetic = opt__magnetic _m.opt.magnetic = opt__magnetic
_m.opt.run_collision_detection = opt__run_collision_detection _m.opt.run_collision_detection = opt__run_collision_detection
@@ -3739,6 +3744,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.opt.iterations, m.opt.iterations,
m.opt.ls_iterations, m.opt.ls_iterations,
m.opt._impl.ls_parallel, m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
m.opt.ls_tolerance, m.opt.ls_tolerance,
m.opt.magnetic, m.opt.magnetic,
m.opt._impl.run_collision_detection, m.opt._impl.run_collision_detection,
mjx/mujoco/mjx/warp/smooth.py
-3
View File
@@ -306,7 +306,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _tendon_shim( def _tendon_shim(
# Model # Model
@@ -488,8 +487,6 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable @ffi.marshal_jax_warp_callable
def tendon(m: types.Model, d: types.Data): def tendon(m: types.Model, d: types.Data):
return _tendon_jax_impl(m, d) return _tendon_jax_impl(m, d)
@tendon.def_vmap @tendon.def_vmap
@ffi.marshal_custom_vmap @ffi.marshal_custom_vmap
def tendon_vmap(unused_axis_size, is_batched, m, d): def tendon_vmap(unused_axis_size, is_batched, m, d):
mjx/mujoco/mjx/warp/types.py
+5
View File
@@ -95,6 +95,7 @@ class OptionWarp(PyTreeNode):
has_fluid: bool has_fluid: bool
is_sparse: bool is_sparse: bool
ls_parallel: bool ls_parallel: bool
ls_parallel_min_step: float
run_collision_detection: bool run_collision_detection: bool
sdf_initpoints: int sdf_initpoints: int
sdf_iterations: int sdf_iterations: int
@@ -885,6 +886,7 @@ _NDIM = {
'opt__iterations': 0, 'opt__iterations': 0,
'opt__ls_iterations': 0, 'opt__ls_iterations': 0,
'opt__ls_parallel': 0, 'opt__ls_parallel': 0,
'opt__ls_parallel_min_step': 0,
'opt__ls_tolerance': 1, 'opt__ls_tolerance': 1,
'opt__magnetic': 2, 'opt__magnetic': 2,
'opt__run_collision_detection': 0, 'opt__run_collision_detection': 0,
@@ -1002,6 +1004,7 @@ _NDIM = {
'iterations': 0, 'iterations': 0,
'ls_iterations': 0, 'ls_iterations': 0,
'ls_parallel': 0, 'ls_parallel': 0,
'ls_parallel_min_step': 0,
'ls_tolerance': 1, 'ls_tolerance': 1,
'magnetic': 2, 'magnetic': 2,
'run_collision_detection': 0, 'run_collision_detection': 0,
@@ -1470,6 +1473,7 @@ _BATCH_DIM = {
'opt__iterations': False, 'opt__iterations': False,
'opt__ls_iterations': False, 'opt__ls_iterations': False,
'opt__ls_parallel': False, 'opt__ls_parallel': False,
'opt__ls_parallel_min_step': False,
'opt__ls_tolerance': True, 'opt__ls_tolerance': True,
'opt__magnetic': True, 'opt__magnetic': True,
'opt__run_collision_detection': False, 'opt__run_collision_detection': False,
@@ -1587,6 +1591,7 @@ _BATCH_DIM = {
'iterations': False, 'iterations': False,
'ls_iterations': False, 'ls_iterations': False,
'ls_parallel': False, 'ls_parallel': False,
'ls_parallel_min_step': False,
'ls_tolerance': True, 'ls_tolerance': True,
'magnetic': True, 'magnetic': True,
'run_collision_detection': False, 'run_collision_detection': False,