chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 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 GeomType
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 kernel as nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -248,54 +250,56 @@ def _obb_filter(
return True
@wp.func
def _broadphase_filter(
# Model:
opt_broadphase_filter: int,
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
geom1: int,
geom2: int,
worldid: int,
) -> bool:
# 1: plane
# 2: sphere
# 4: aabb
# 8: obb
def _broadphase_filter(opt_broadphase_filter: int):
@wp.func
def func(
# Model:
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
geom1: int,
geom2: int,
worldid: int,
) -> bool:
# 1: plane
# 2: sphere
# 4: aabb
# 8: obb
center1 = geom_aabb[geom1, 0]
center2 = geom_aabb[geom2, 0]
size1 = geom_aabb[geom1, 1]
size2 = geom_aabb[geom2, 1]
rbound1 = geom_rbound[worldid, geom1]
rbound2 = geom_rbound[worldid, geom2]
margin1 = geom_margin[worldid, geom1]
margin2 = geom_margin[worldid, geom2]
xpos1 = geom_xpos_in[worldid, geom1]
xpos2 = geom_xpos_in[worldid, geom2]
xmat1 = geom_xmat_in[worldid, geom1]
xmat2 = geom_xmat_in[worldid, geom2]
center1 = geom_aabb[geom1, 0]
center2 = geom_aabb[geom2, 0]
size1 = geom_aabb[geom1, 1]
size2 = geom_aabb[geom2, 1]
rbound1 = geom_rbound[worldid, geom1]
rbound2 = geom_rbound[worldid, geom2]
margin1 = geom_margin[worldid, geom1]
margin2 = geom_margin[worldid, geom2]
xpos1 = geom_xpos_in[worldid, geom1]
xpos2 = geom_xpos_in[worldid, geom2]
xmat1 = geom_xmat_in[worldid, geom1]
xmat2 = geom_xmat_in[worldid, geom2]
if rbound1 == 0.0 or rbound2 == 0.0:
if opt_broadphase_filter & int(BroadphaseFilter.PLANE.value):
return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
else:
if opt_broadphase_filter & int(BroadphaseFilter.SPHERE.value):
if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
return False
if opt_broadphase_filter & int(BroadphaseFilter.AABB.value):
if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
if opt_broadphase_filter & int(BroadphaseFilter.OBB.value):
if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
if rbound1 == 0.0 or rbound2 == 0.0:
if wp.static(opt_broadphase_filter & int(BroadphaseFilter.PLANE.value)):
return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
else:
if wp.static(opt_broadphase_filter & int(BroadphaseFilter.SPHERE.value)):
if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
return False
if wp.static(opt_broadphase_filter & int(BroadphaseFilter.AABB.value)):
if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
if wp.static(opt_broadphase_filter & int(BroadphaseFilter.OBB.value)):
if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
return True
return True
return func
@wp.func
@@ -413,82 +417,81 @@ def _sap_range(
sap_range_out[worldid, geomid] = limit - geomid
@wp.kernel
def _sap_broadphase(
# Model:
ngeom: int,
opt_broadphase_filter: int,
geom_type: wp.array(dtype=int),
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
nxn_pairid: wp.array(dtype=int),
# Data in:
nworld_in: int,
nconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore
sap_cumulative_sum_in: wp.array(dtype=int), # kernel_analyzer: ignore
# In:
nsweep_in: int,
# Data out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_hftri_index_out: wp.array(dtype=int),
collision_pairid_out: wp.array(dtype=int),
collision_worldid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
worldgeomid = wp.tid()
@cache_kernel
def _sap_broadphase(broadphase_filter):
@nested_kernel
def kernel(
# Model:
ngeom: int,
geom_type: wp.array(dtype=int),
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
nxn_pairid: wp.array(dtype=int),
# Data in:
nworld_in: int,
nconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore
sap_cumulative_sum_in: wp.array(dtype=int), # kernel_analyzer: ignore
# In:
nsweep_in: int,
# Data out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_hftri_index_out: wp.array(dtype=int),
collision_pairid_out: wp.array(dtype=int),
collision_worldid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
worldgeomid = wp.tid()
nworldgeom = nworld_in * ngeom
nworkpackages = sap_cumulative_sum_in[nworldgeom - 1]
nworldgeom = nworld_in * ngeom
nworkpackages = sap_cumulative_sum_in[nworldgeom - 1]
while worldgeomid < nworkpackages:
# binary search to find current and next geom pair indices
i = _binary_search(sap_cumulative_sum_in, worldgeomid, 0, nworldgeom)
j = i + worldgeomid + 1
while worldgeomid < nworkpackages:
# binary search to find current and next geom pair indices
i = _binary_search(sap_cumulative_sum_in, worldgeomid, 0, nworldgeom)
j = i + worldgeomid + 1
if i > 0:
j -= sap_cumulative_sum_in[i - 1]
if i > 0:
j -= sap_cumulative_sum_in[i - 1]
worldid = i // ngeom
i = i % ngeom
j = j % ngeom
worldid = i // ngeom
i = i % ngeom
j = j % ngeom
# get geom indices and swap if necessary
geom1 = sap_sort_index_in[worldid, i]
geom2 = sap_sort_index_in[worldid, j]
# get geom indices and swap if necessary
geom1 = sap_sort_index_in[worldid, i]
geom2 = sap_sort_index_in[worldid, j]
# find linear index of (geom1, geom2) in upper triangular nxn_pairid
if geom2 < geom1:
idx = upper_tri_index(ngeom, geom2, geom1)
else:
idx = upper_tri_index(ngeom, geom1, geom2)
# find linear index of (geom1, geom2) in upper triangular nxn_pairid
if geom2 < geom1:
idx = upper_tri_index(ngeom, geom2, geom1)
else:
idx = upper_tri_index(ngeom, geom1, geom2)
if nxn_pairid[idx] < -1:
worldgeomid += nsweep_in
continue
if nxn_pairid[idx] < -1:
continue
if _broadphase_filter(
opt_broadphase_filter, geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
):
_add_geom_pair(
geom_type,
nxn_pairid,
nconmax_in,
geom1,
geom2,
worldid,
idx,
collision_pair_out,
collision_hftri_index_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
if broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid):
_add_geom_pair(
geom_type,
nxn_pairid,
nconmax_in,
geom1,
geom2,
worldid,
idx,
collision_pair_out,
collision_hftri_index_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
worldgeomid += nsweep_in
return kernel
def _segmented_sort(tile_size: int):
@@ -591,12 +594,12 @@ def sap_broadphase(m: Model, d: Data):
# estimate number of overlap checks
# assumes each geom has 5 other geoms (batched over all worlds)
nsweep = 5 * nworldgeom
broadphase_filter = _broadphase_filter(m.opt.broadphase_filter)
wp.launch(
kernel=_sap_broadphase,
kernel=_sap_broadphase(broadphase_filter),
dim=nsweep,
inputs=[
m.ngeom,
m.opt.broadphase_filter,
m.geom_type,
m.geom_aabb,
m.geom_rbound,
@@ -620,50 +623,51 @@ def sap_broadphase(m: Model, d: Data):
)
@wp.kernel
def _nxn_broadphase(
# Model:
opt_broadphase_filter: int,
geom_type: wp.array(dtype=int),
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
nxn_geom_pair: wp.array(dtype=wp.vec2i),
nxn_pairid: wp.array(dtype=int),
# Data in:
nconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# Data out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_hftri_index_out: wp.array(dtype=int),
collision_pairid_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
@cache_kernel
def _nxn_broadphase(broadphase_filter):
@nested_kernel
def kernel(
# Model:
geom_type: wp.array(dtype=int),
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
nxn_geom_pair: wp.array(dtype=wp.vec2i),
nxn_pairid: wp.array(dtype=int),
# Data in:
nconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# Data out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_hftri_index_out: wp.array(dtype=int),
collision_pairid_out: wp.array(dtype=int),
collision_worldid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
_add_geom_pair(
geom_type,
nxn_pairid,
nconmax_in,
geom1,
geom2,
worldid,
elementid,
collision_pair_out,
collision_hftri_index_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
worldid, elementid = wp.tid()
geom = nxn_geom_pair[elementid]
geom1 = geom[0]
geom2 = geom[1]
if broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid):
_add_geom_pair(
geom_type,
nxn_pairid,
nconmax_in,
geom1,
geom2,
worldid,
elementid,
collision_pair_out,
collision_hftri_index_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
return kernel
@event_scope
@@ -681,11 +685,11 @@ def nxn_broadphase(m: Model, d: Data):
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
"""
broadphase_filter = _broadphase_filter(m.opt.broadphase_filter)
wp.launch(
_nxn_broadphase,
_nxn_broadphase(broadphase_filter),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[
m.opt.broadphase_filter,
m.geom_type,
m.geom_aabb,
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)
mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
# box-box currently supports up to 8 contacts
MULTI_CONTACT_COUNT = 8
mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
@@ -94,7 +95,7 @@ class SupportPoint:
@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 (
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,
):
"""Find distance within a tolerance between two geoms."""
is_discrete = discrete_geoms(geomtype1, geomtype2)
is_discrete = _discrete_geoms(geomtype1, geomtype2)
cutoff2 = cutoff * cutoff
simplex = mat43()
simplex1 = mat43()
@@ -1203,7 +1204,7 @@ def _polytope4(
@wp.func
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."""
is_discrete = discrete_geoms(geomtype1, geomtype2)
is_discrete = _discrete_geoms(geomtype1, geomtype2)
upper = FLOAT_MAX
upper2 = FLOAT_MAX
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
@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
@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]]
v2i = vert_index[face[1]]
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
@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()
for i in range(len1):
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
# return 1 if found
@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()
for i in range(nface):
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
@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)
res = wp.vec2i()
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
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)
res = wp.vec2i()
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
@wp.func
def mesh_normals(
def _mesh_normals(
# In:
feature_dim: int,
feature_index: wp.vec3i,
@@ -1420,10 +1472,10 @@ def mesh_normals(
v3_num = polymapnum[vertadr + v3]
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:
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:
return 0, normals, indices
@@ -1440,7 +1492,7 @@ def mesh_normals(
v2_num = polymapnum[vertadr + v2]
# 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:
return 0, normals, indices
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
@wp.func
def mesh_edge_normals(
def _mesh_edge_normals(
# In:
dim: int,
mat: wp.mat33,
@@ -1511,7 +1563,7 @@ def mesh_edge_normals(
# try recovering box normal from collision normal
@wp.func
def box_normals2(mat: wp.mat33, n: wp.vec3):
def _box_normals2(mat: wp.mat33, n: wp.vec3):
normals = polyverts()
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
@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()
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
if c == 1:
return 1, normals, indices
return box_normals2(mat, dir)
return _box_normals2(mat, dir)
if feature_dim == 2:
c = 0
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
if c == 2:
return 2, normals, indices
return box_normals2(mat, dir)
return _box_normals2(mat, dir)
if feature_dim == 1:
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
@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()
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
@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()
# 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
@wp.func
def mesh_face(
def _mesh_face(
# In:
mat: wp.mat33,
pos: wp.vec3,
@@ -1706,19 +1758,19 @@ def mesh_face(
@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
res = wp.cross(v2 - v1, v3 - v1)
return wp.dot(res, v1), res
@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
@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()
ab = b - a
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
@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()
witness2 = mat3c()
@@ -1746,10 +1798,10 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
pn = polyverts()
pd = polyvec()
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
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
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])
# determine if P and Q are in the halfspace of the clipping edge
inside1 = halfspace(face1[e], pn[e], P)
inside2 = halfspace(face1[e], pn[e], Q)
inside1 = _halfspace(face1[e], pn[e], P)
inside2 = _halfspace(face1[e], pn[e], Q)
# PQ entirely outside the clipping edge, skip
if not inside1 and not inside2:
@@ -1784,7 +1836,7 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
continue
# 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:
clipped[nclipped] = res
nclipped += 1
@@ -1804,6 +1856,13 @@ def polygon_clip(face1: polyverts, nface1: int, face2: polyverts, nface2: int, n
if npolygon < 1:
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
for i in range(npolygon):
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
@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
):
witness1 = mat3c()
@@ -1845,40 +1904,40 @@ def multicontact(
polymap = geom2.mesh_polymap
# get dimensions of features of geoms 1 and 2
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)
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)
dir = x2 - x1
dir_neg = -dir
# get all possible face normals for each geom
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):
nnorms1, n1, idx1 = mesh_normals(
nnorms1, n1, idx1 = _mesh_normals(
nface1, feature_index1, geom1.rot, geom1.vertadr, geom1.mesh_polyadr, polynormal, polymapadr, polymapnum, polymap
)
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):
nnorms2, n2, idx2 = mesh_normals(
nnorms2, n2, idx2 = _mesh_normals(
nface2, feature_index2, geom2.rot, geom2.vertadr, geom2.mesh_polyadr, polynormal, polymapadr, polymapnum, polymap
)
# determine if any two face normals match
is_edge_contact_geom1 = 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:
# check if edge-face collision
if nface1 < 3 and nface1 <= nface2:
nnorms1 = 0
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]
)
elif geomtype1 == int(GeomType.MESH.value):
nnorms1, n1, endverts = mesh_edge_normals(
nnorms1, n1, endverts = _mesh_edge_normals(
nface1,
geom1.rot,
geom1.pos,
@@ -1895,7 +1954,7 @@ def multicontact(
feature_vertex1[1],
feature_index1[0],
)
nres, res = aligned_face_edge(n1, nnorms1, n2, nnorms2)
nres, res = _aligned_face_edge(n1, nnorms1, n2, nnorms2)
if not nres:
return 1, witness1, witness2
is_edge_contact_geom1 = 1
@@ -1904,11 +1963,11 @@ def multicontact(
elif nface2 < 3:
nnorms2 = 0
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]
)
elif geomtype2 == int(GeomType.MESH.value):
nnorms2, n2, endverts = mesh_edge_normals(
nnorms2, n2, endverts = _mesh_edge_normals(
nface2,
geom2.rot,
geom2.pos,
@@ -1925,7 +1984,7 @@ def multicontact(
feature_vertex2[1],
feature_index2[0],
)
nres, res = aligned_face_edge(n2, nnorms2, n1, nnorms1)
nres, res = _aligned_face_edge(n2, nnorms2, n1, nnorms1)
if not nres:
return 1, witness1, witness2
is_edge_contact_geom2 = 1
@@ -1944,9 +2003,9 @@ def multicontact(
else:
ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i])
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):
nface1, face1 = mesh_face(
nface1, face1 = _mesh_face(
geom1.rot, geom1.pos, geom1.vertadr, geom1.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, ind
)
@@ -1957,9 +2016,9 @@ def multicontact(
nface2 = 2
else:
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):
nface2, face2 = mesh_face(
nface2, face2 = _mesh_face(
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
if is_edge_contact_geom1:
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
if is_edge_contact_geom2:
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
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
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)
if margin1 > 0.0:
x1 += margin1 * n
@@ -2021,9 +2080,9 @@ def ccd(
face_map: wp.array(dtype=int),
horizon: wp.array(dtype=int),
):
"""General convex collision detection via GJK/EPA."""
witness1 = mat3c()
witness2 = mat3c()
"""General convex collision detection via GJK/EPA."""
margin1 = 0.0
margin2 = 0.0
@@ -2044,7 +2103,7 @@ def ccd(
witness1[0] = result.x1
witness2[0] = result.x2
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
witness2[0] = x2
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 (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:
return dist, num, w1, w2
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)
self.assertEqual(ncon, 5)
self.assertEqual(ncon, 4)
def test_box_box_ccd2(self):
"""Test box-box multiccd 2."""
@@ -441,7 +441,7 @@ class GJKTest(absltest.TestCase):
)
_, ncon, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=True)
self.assertEqual(ncon, 5)
self.assertEqual(ncon, 4)
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
@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
def deriv_smooth_vel(m: Model, d: Data, flg_forward: bool = True):
"""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:
wp.launch(
_qfrc_forward,
dim=(d.nworld, m.nv),
inputs=[d.qfrc_smooth, d.qfrc_constraint],
outputs=[d.qfrc_integration],
)
wp.copy(d.qfrc_integration, d.efc.Ma)
else:
# qfrc = qM @ qacc
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),
dof_Madr: wp.array(dtype=int),
dof_damping: wp.array2d(dtype=float),
# 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),
qM_integration_out: wp.array3d(dtype=float),
):
worldid, tid = wp.tid()
@@ -331,7 +327,6 @@ def _euler_damp_qfrc_sparse(
adr = dof_Madr[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):
@@ -343,11 +338,8 @@ def _euler_sparse(m: Model, d: Data):
m.opt.timestep,
m.dof_Madr,
m.dof_damping,
d.qfrc_smooth,
d.qfrc_constraint,
],
outputs=[
d.qfrc_integration,
d.qM_integration,
],
)
@@ -358,7 +350,7 @@ def _euler_sparse(m: Model, d: Data):
d.qLD_integration,
d.qLDiagInv_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),
# Data in:
qM_in: wp.array3d(dtype=float),
qfrc_smooth_in: wp.array2d(dtype=float),
qfrc_constraint_in: wp.array2d(dtype=float),
efc_Ma_in: wp.array2d(dtype=float),
# In:
adr_in: wp.array(dtype=int),
# Data out:
@@ -388,14 +379,10 @@ def _tile_euler_dense(tile: TileSet):
damping_scaled = damping_tile * timestep
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,))
qfrc_constraint_tile = wp.tile_load(qfrc_constraint_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
qfrc_tile = qfrc_smooth_tile + qfrc_constraint_tile
Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_cholesky(qm_integration_tile)
qacc_tile = wp.tile_cholesky_solve(L_tile, qfrc_tile)
wp.tile_store(qacc_integration_out[worldid], qacc_tile, offset=(dofid))
qacc_integration_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
wp.tile_store(qacc_integration_out[worldid], qacc_integration_tile, offset=(dofid))
return euler_dense
@@ -413,7 +400,7 @@ def euler(m: Model, d: Data):
wp.launch_tiled(
_tile_euler_dense(tile),
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],
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.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 IntegratorType
@@ -133,14 +132,13 @@ class ForwardTest(parameterized.TestCase):
mjm, mjd, _, _ = test_util.fixture("pendula.xml", kick=True, eulerdamp=eulerdamp, sparse=sparse)
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)
d = mjwarp.put_data(mjm, mjd)
mujoco.mj_forward(mjm, mjd)
mujoco.mj_Euler(mjm, mjd)
mjwarp.solve(m, d) # compute efc.Ma
mjwarp.euler(m, d)
_assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
@@ -220,9 +218,11 @@ class ForwardTest(parameterized.TestCase):
m = mjwarp.put_model(mjm)
d = mjwarp.put_data(mjm, mjd)
mjwarp.implicit(m, d)
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.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)
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")
@@ -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)
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")
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.")
# reduction
if (mjm.sensor_intprm[is_contact_sensor, 1] != 1).any():
raise NotImplementedError(f"Contact sensor: only mindist reduction is implemented.")
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 and maxforce reduction are implemented.")
# TODO(team): remove after _update_gradient for Newton uses tile operations for islands
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),
is_sparse=bool(is_sparse),
ls_parallel=False,
ls_parallel_min_step=1.0e-6, # TODO(team): determine good default setting
gjk_iterations=MJ_CCD_ITERATIONS,
epa_iterations=MJ_CCD_ITERATIONS,
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 geom2="plane"/>',
'<contact site="site"/>',
'<contact reduce="maxforce"/>',
'<contact reduce="netforce"/>',
'<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:
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))
for i in range(nvert):
for j in range(nvert):
for x in range(3):
force[i, x] -= flex_bending[17*edgeid + 4*i + j] * flexvert_xpos_in[worldid, v[j]][x]
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 + 16] * frc[i, x]
for i in range(nvert):
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)
y0 = wp.max(0.0, wp.min(float(nrow - 2), wp.floor(y)))
if i == 1:
z0 = hfield_data[adr + int(wp.round(y0)) * nrow + ncol - 1]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * 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)) * ncol + ncol - 1]
else:
z0 = hfield_data[adr + int(wp.round(y0)) * nrow]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * nrow]
z0 = hfield_data[adr + int(wp.round(y0 + 0.0)) * ncol]
z1 = hfield_data[adr + int(wp.round(y0 + 1.0)) * ncol]
# side normal to y-axis
else:
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)))
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]
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))]
# 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
def _contact_match(
# Model:
opt_cone: int,
sensor_objid: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_intprm: wp.array2d(dtype=int),
sensor_contact_adr: wp.array(dtype=int),
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
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_efc_address_in: wp.array2d(dtype=int),
contact_worldid_in: wp.array(dtype=int),
efc_force_in: wp.array2d(dtype=float),
# Data out:
sensor_contact_nmatch_out: wp.array2d(dtype=int),
sensor_contact_matchid_out: wp.array3d(dtype=int),
@@ -2009,6 +2017,7 @@ def _contact_match(
# sensor information
objid = sensor_objid[sensorid]
refid = sensor_refid[sensorid]
reduce = sensor_intprm[sensorid, 1]
# contact information
geom = contact_geom_in[contactid]
@@ -2022,8 +2031,27 @@ def _contact_match(
contactmatchid = wp.atomic_add(sensor_contact_nmatch_out[worldid], contactsensorid, 1)
sensor_contact_matchid_out[worldid, contactsensorid, contactmatchid] = contactid
# TODO(thowell): alternative criteria
sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = contact_dist_in[contactid]
if reduce == 1: # mindist
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
if geom1geom0:
@@ -2155,20 +2183,28 @@ def sensor_acc(m: Model, d: Data):
if m.sensor_contact_adr.size:
# match criteria
d.sensor_contact_nmatch.zero_()
d.sensor_contact_matchid.zero_()
d.sensor_contact_criteria.zero_()
d.sensor_contact_matchid.fill_(-1)
d.sensor_contact_criteria.fill_(1.0e32)
wp.launch(
_contact_match,
dim=(m.sensor_contact_adr.size, d.nconmax),
inputs=[
m.opt.cone,
m.sensor_objid,
m.sensor_refid,
m.sensor_intprm,
m.sensor_contact_adr,
d.njmax,
d.ncon,
d.contact.dist,
d.contact.frame,
d.contact.friction,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
d.efc.force,
],
outputs=[
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(
'type="sphere" size=".1"',
'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):
"""Test contact sensor."""
@@ -419,22 +419,18 @@ class SensorTest(parameterized.TestCase):
contact_sensor = ""
# data combinations
field = ["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)
datas = ["found", "force dist normal", "torque pos tangent", "found force torque dist pos normal tangent"]
for num in [1, 2, 3, 4, 5]:
for geoms in [
'geom1="plane" geom2="geom"',
'geom1="geom" geom2="plane"',
'geom1="plane" geom2="sphere"',
'geom1="sphere" geom2="plane"',
'geom1="geom" geom2="sphere"',
'geom1="sphere" geom2="geom"',
]:
for data in datas:
data = " ".join(data)
contact_sensor += f'<contact {geoms} num="{num}" reduce="mindist" data="{data}"/>'
for geoms in [
'geom1="plane" geom2="geom"',
'geom1="geom" geom2="plane"',
'geom1="sphere" geom2="plane"',
'geom1="geom" geom2="sphere"',
]:
for num in [1, 3, 5]:
for reduce in ["mindist", "maxforce"]:
for data in datas:
contact_sensor += f'<contact {geoms} num="{num}" reduce="{reduce}" data="{data}"/>'
_MJCF = f"""
<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
def linesearch_parallel_fused(
# Model:
nlsp: int,
opt_impratio: wp.array(dtype=float),
opt_ls_parallel_min_step: float,
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
@@ -750,7 +757,7 @@ def linesearch_parallel_fused(
if efc_done_in[worldid]:
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)
@@ -846,6 +853,7 @@ def linesearch_parallel_fused(
def linesearch_parallel_best_alpha(
# Model:
nlsp: int,
opt_ls_parallel_min_step: float,
# Data in:
efc_done_in: wp.array(dtype=bool),
efc_cost_candidate_in: wp.array2d(dtype=float),
@@ -867,7 +875,7 @@ def linesearch_parallel_best_alpha(
best_cost = cost
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):
@@ -877,6 +885,7 @@ def _linesearch_parallel(m: types.Model, d: types.Data):
inputs=[
m.nlsp,
m.opt.impratio,
m.opt.ls_parallel_min_step,
d.njmax,
d.ncon,
d.ne,
@@ -900,7 +909,7 @@ def _linesearch_parallel(m: types.Model, d: types.Data):
wp.launch(
linesearch_parallel_best_alpha,
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],
)
@@ -1589,26 +1598,6 @@ def update_gradient_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
def update_gradient_set_h_qM_lower_sparse(
# Model:
@@ -1627,32 +1616,11 @@ def update_gradient_set_h_qM_lower_sparse(
i = qM_fullm_i[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
def update_gradient_copy_lower_triangle(
# 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(
def update_gradient_JTDAJ_sparse(
# Data in:
njmax_in: int,
nefc_in: wp.array(dtype=int),
@@ -1673,12 +1641,14 @@ def update_gradient_JTDAJ(
dofi = (int(sqrt(float(1 + 8 * elementid))) - 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]
# TODO(team): sparse efc_J
efc_Ji = efc_J_in[worldid, 0, dofi]
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):
if efc_state == int(types.ConstraintState.QUADRATIC.value) and efc_D != 0.0:
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_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
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 ?
@@ -1904,11 +1925,19 @@ def _update_gradient(m: types.Model, d: types.Data):
smooth.solve_m(m, d, d.efc.Mgrad, d.efc.grad)
elif m.opt.solver == types.SolverType.NEWTON:
# 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:
wp.launch(
update_gradient_zero_h_lower,
dim=(d.nworld, m.dof_tri_row.size),
inputs=[m.dof_tri_row, m.dof_tri_col, d.efc.done],
update_gradient_JTDAJ_sparse,
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],
)
wp.launch(
@@ -1919,28 +1948,20 @@ def _update_gradient(m: types.Model, d: types.Data):
)
else:
wp.launch(
update_gradient_copy_lower_triangle,
dim=(d.nworld, m.dof_tri_row.size),
inputs=[m.dof_tri_row, m.dof_tri_col, d.qM, d.efc.done],
update_gradient_JTDAJ_dense,
dim=(d.nworld, lower_triangle_dim),
inputs=[
d.njmax,
d.nefc,
d.qM,
d.efc.J,
d.efc.D,
d.efc.state,
d.efc.done,
],
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:
# 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
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):
@parameterized.product(cone=tuple(ConeType), solver_=tuple(SolverType))
def test_cost(self, cone, solver_):
"""Tests cost function is correct."""
def test_constraint_update(self, cone, solver_):
"""Tests _update_constraint function is correct."""
for keyframe in range(3):
mjm, mjd, m, d = test_util.fixture(
"constraints.xml",
@@ -60,14 +60,37 @@ class SolverTest(parameterized.TestCase):
mujoco.mj_constraintUpdate(mjm, mjd, jaref - mjd.efc_aref, cost, 0)
return cost
mj_cost = cost(mjd.qacc)
mjd_cost = cost(mjd.qacc)
# 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_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):
"""Test linesearch initialization."""
@@ -160,10 +183,6 @@ class SolverTest(parameterized.TestCase):
def test_parallel_linesearch(self, cone):
"""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(
"humanoid/humanoid.xml",
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
def _apply_ft(
# Model:
@@ -232,6 +193,9 @@ def _apply_ft(
accumul = float(0.0)
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
parentid = bodyid
while parentid != 0 and parentid != dofbodyid:
@@ -240,7 +204,6 @@ def _apply_ft(
continue # body is not part of the subtree
offset = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]]
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))
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,
measure_alloc: 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.
Args:
@@ -217,22 +217,20 @@ def benchmark(
list: Number of contacts.
list: Number of constraints.
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 = {}
ncon, nefc, solver_niter = [], [], []
with warp_util.EventTracer(enabled=event_trace) as tracer:
# capture the whole function as a CUDA graph
jit_beg = time.perf_counter()
with wp.ScopedCapture() as capture:
fn(m, d)
jit_end = time.perf_counter()
jit_duration = jit_end - jit_beg
graph = capture.graph
time_vec = np.zeros(nstep)
@@ -258,14 +256,15 @@ def benchmark(
else:
trace = tracer.trace()
if measure_alloc:
ncon.append(d.ncon.numpy()[0])
nefc.append(np.sum(d.nefc.numpy()))
ncon.append(np.max([d.ncon.numpy()[0], d.ncollision.numpy()[0]]))
nefc.append(np.max(d.nefc.numpy()))
if measure_solver_niter:
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)
return jit_duration, run_duration, trace, ncon, nefc, solver_niter
return jit_duration, run_duration, trace, ncon, nefc, solver_niter, nsuccess
class BenchmarkSuite:
@@ -322,17 +321,20 @@ class BenchmarkSuite:
mujoco.mj_forward(mjm, mjd)
wp.init()
if os.environ.get("ASV_CACHE_KERNELS", "false").lower() == "false":
wp.clear_kernel_cache()
free_before = wp.get_device().free_memory
m = io.put_model(mjm)
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 = {
"jit_duration": jit_duration,
"solver_niter_mean": np.mean(solver_niter),
"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=""):
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
epa_iterations: number of Epa iterations in the convex narrowphase
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)
has_fluid: True if wind, density, or viscosity are non-zero at put_model time
density: density of medium
@@ -587,7 +588,8 @@ class Option:
is_sparse: bool
gjk_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)
has_fluid: bool
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.")
_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")
_KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
_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:
try:
val = str(enum_fields[key][val.upper()])
val = int(enum_fields[key][val.upper()])
except KeyError:
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}")
if i < len(attrs) - 1:
obj = getattr(obj, attr)
else:
elif key not in enum_fields:
try:
val = type(getattr(obj, attr))(ast.literal_eval(val))
except (SyntaxError, ValueError):
raise app.UsageError(f"Unrecognized value for field: {key}")
setattr(obj, attr, val)
setattr(obj, attr, val)
def _compile_step(m, d):
@@ -138,6 +138,8 @@ def _main(argv: Sequence[str]) -> None:
mjm = _load_model(epath.Path(argv[1]))
mjd = mujoco.MjData(mjm)
if mjm.nkey > 0 and _KEYFRAME.value > -1:
mujoco.mj_resetDataKeyframe(mjm, mjd, _KEYFRAME.value)
mujoco.mj_forward(mjm, mjd)
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__ls_iterations: int,
opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
opt__ls_tolerance: wp.array(dtype=float),
opt__magnetic: wp.array(dtype=wp.vec3),
opt__run_collision_detection: bool,
@@ -734,6 +735,7 @@ def _forward_shim(
_m.opt.iterations = opt__iterations
_m.opt.ls_iterations = opt__ls_iterations
_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.magnetic = opt__magnetic
_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.ls_iterations,
m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
m.opt.ls_tolerance,
m.opt.magnetic,
m.opt._impl.run_collision_detection,
@@ -2358,6 +2361,7 @@ def _step_shim(
opt__iterations: int,
opt__ls_iterations: int,
opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
opt__ls_tolerance: wp.array(dtype=float),
opt__magnetic: wp.array(dtype=wp.vec3),
opt__run_collision_detection: bool,
@@ -2773,6 +2777,7 @@ def _step_shim(
_m.opt.iterations = opt__iterations
_m.opt.ls_iterations = opt__ls_iterations
_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.magnetic = opt__magnetic
_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.ls_iterations,
m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
m.opt.ls_tolerance,
m.opt.magnetic,
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}
)
@ffi.format_args_for_warp
def _tendon_shim(
# Model
@@ -488,8 +487,6 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable
def tendon(m: types.Model, d: types.Data):
return _tendon_jax_impl(m, d)
@tendon.def_vmap
@ffi.marshal_custom_vmap
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
is_sparse: bool
ls_parallel: bool
ls_parallel_min_step: float
run_collision_detection: bool
sdf_initpoints: int
sdf_iterations: int
@@ -885,6 +886,7 @@ _NDIM = {
'opt__iterations': 0,
'opt__ls_iterations': 0,
'opt__ls_parallel': 0,
'opt__ls_parallel_min_step': 0,
'opt__ls_tolerance': 1,
'opt__magnetic': 2,
'opt__run_collision_detection': 0,
@@ -1002,6 +1004,7 @@ _NDIM = {
'iterations': 0,
'ls_iterations': 0,
'ls_parallel': 0,
'ls_parallel_min_step': 0,
'ls_tolerance': 1,
'magnetic': 2,
'run_collision_detection': 0,
@@ -1470,6 +1473,7 @@ _BATCH_DIM = {
'opt__iterations': False,
'opt__ls_iterations': False,
'opt__ls_parallel': False,
'opt__ls_parallel_min_step': False,
'opt__ls_tolerance': True,
'opt__magnetic': True,
'opt__run_collision_detection': False,
@@ -1587,6 +1591,7 @@ _BATCH_DIM = {
'iterations': False,
'ls_iterations': False,
'ls_parallel': False,
'ls_parallel_min_step': False,
'ls_tolerance': True,
'magnetic': True,
'run_collision_detection': False,