From bf04103b6fc0b68b78ef53a3999f7dac7e170cd7 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 19 Aug 2025 16:00:39 -0700 Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub. #MIGRATION_3P_PY_MUJOCO_WARP__DEFAULT - 19dde09e179dbbaf4d07084185074f6a774b3255 update derivative.py by Taylor Howell - 04bc115c3766e974ef000cca16d40bacb92c01d2 fix _ray_hfield by Taylor Howell - 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 - 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 - 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 - 4f731941cb9dea397d4f71cf8151100b83eae576 A few airspeed velocity improvements. (#626) by Erik Frey PiperOrigin-RevId: 797054114 Change-Id: I380e0b9446f94202b2389e2a760d6de5f0ec5af7 --- .../mujoco_warp/_src/collision_driver.py | 320 +++++++++--------- .../mujoco_warp/_src/collision_gjk.py | 167 ++++++--- .../mujoco_warp/_src/collision_gjk_test.py | 4 +- .../mujoco_warp/_src/derivative.py | 19 +- .../third_party/mujoco_warp/_src/forward.py | 25 +- .../mujoco_warp/_src/forward_test.py | 16 +- .../mjx/third_party/mujoco_warp/_src/io.py | 5 +- .../third_party/mujoco_warp/_src/io_test.py | 1 - .../third_party/mujoco_warp/_src/passive.py | 18 +- .../mjx/third_party/mujoco_warp/_src/ray.py | 12 +- .../third_party/mujoco_warp/_src/sensor.py | 44 ++- .../mujoco_warp/_src/sensor_test.py | 28 +- .../third_party/mujoco_warp/_src/solver.py | 163 +++++---- .../mujoco_warp/_src/solver_test.py | 37 +- .../third_party/mujoco_warp/_src/support.py | 43 +-- .../third_party/mujoco_warp/_src/test_util.py | 28 +- .../mjx/third_party/mujoco_warp/_src/types.py | 4 +- .../mjx/third_party/mujoco_warp/viewer.py | 10 +- mjx/mujoco/mjx/warp/forward.py | 6 + mjx/mujoco/mjx/warp/smooth.py | 3 - mjx/mujoco/mjx/warp/types.py | 5 + 21 files changed, 528 insertions(+), 430 deletions(-) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py index 4cdd85c2..199c0b62 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py @@ -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 `` 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, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py index ff999d2d..089d5856 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py @@ -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 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py index f47f7e01..64dd7945 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py @@ -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__": diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py index dcd70f4b..ad45b7c4 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py @@ -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) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py index 3ae65b12..31309249 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -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, ) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py index fa428f46..73087c3c 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py @@ -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") diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index 4ec248bd..c9742982 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -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), diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io_test.py index b197dd9a..2d5932d8 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io_test.py @@ -366,7 +366,6 @@ class IOTest(parameterized.TestCase): '', '', '', - '', '', '', ) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py index da4b4623..a91858f3 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -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]] diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py index 6bc72ce5..787c4cd6 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py @@ -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 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py index bc561af6..3b82a825 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -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, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor_test.py index 853b71c2..96dcdda0 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor_test.py @@ -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'' + 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'' _MJCF = f""" diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py index c77de1e1..7d208219 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -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 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver_test.py index 610a8f68..fa192929 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver_test.py @@ -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, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py index 5d276fbe..1018a6b5 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -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: diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py index b78de634..be951a31 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py @@ -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=""): diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py index 55f09ffe..29d2693f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -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) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py index 0a27a2bd..56b02e0f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py @@ -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: diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 3ba74834..d2437fdc 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -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, diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 3b881891..ce4bd16e 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -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): diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 9c3a188d..8b556e5f 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -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,