diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py index 1b76ba88..238141eb 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py @@ -64,7 +64,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.solver import solve as solve from mujoco.mjx.third_party.mujoco_warp._src.support import contact_force as contact_force from mujoco.mjx.third_party.mujoco_warp._src.support import mul_m as mul_m from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as xfrc_accumulate +from mujoco.mjx.third_party.mujoco_warp._src.test_util import BenchmarkSuite as BenchmarkSuite from mujoco.mjx.third_party.mujoco_warp._src.test_util import benchmark as benchmark +from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter as BroadphaseFilter from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType as BroadphaseType from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType as ConeType from mujoco.mjx.third_party.mujoco_warp._src.types import Constraint as Constraint @@ -72,6 +74,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Contact as Contact from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit as DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import DynType as DynType from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit as EnableBit +from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType as IntegratorType from mujoco.mjx.third_party.mujoco_warp._src.types import JointType as JointType from mujoco.mjx.third_party.mujoco_warp._src.types import Option as Option from mujoco.mjx.third_party.mujoco_warp._src.types import SolverType as SolverType diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py index 8ddb801e..b29f6d83 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py @@ -321,7 +321,8 @@ def ccd_kernel_builder( x1 += hfield_prism_vertex(geom1.hfprism, i) x1 = x1 / 6.0 - dist, x1, x2 = ccd( + dist, count, witness1, witness2 = ccd( + False, 1e-6, 0.0, gjk_iterations, @@ -344,13 +345,13 @@ def ccd_kernel_builder( epa_map_in[tid], epa_horizon_in[tid], ) - count = 0 - if dist < 0.0: - count = 1 - - points[0] = 0.5 * (x1 + x2) - normal = x1 - x2 + if dist >= 0.0: + count = 0 + return + for i in range(count): + points[i] = 0.5 * (witness1[i] + witness2[i]) + normal = witness1[0] - witness2[0] frame = make_frame(normal) for i in range(count): # limit maximum number of contacts with height field diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py index 130121d4..964a673a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py @@ -475,7 +475,7 @@ class CollisionTest(parameterized.TestCase): @classmethod def setUpClass(cls): - register_sdf_plugins(mjwarp._src.collision_sdf) + register_sdf_plugins(mjwarp) @parameterized.parameters(_SDF_SDF.keys()) def test_sdf_collision(self, fixture): 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 8128262a..81cb0748 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 @@ -27,7 +27,22 @@ FLOAT_MIN = -1e30 FLOAT_MAX = 1e30 MJ_MINVAL2 = MJ_MINVAL * MJ_MINVAL +# TODO(kbayes): write out formulas to derive these constants +FACE_TOL = 0.99999872 +EDGE_TOL = 0.00159999931 + +MAX_POLYVERT = 15 +polyverts = wp.types.matrix(shape=(MAX_POLYVERT, 3), dtype=float) +polyclip = wp.types.matrix(shape=(2 * MAX_POLYVERT, 3), dtype=float) +polyvec = wp.types.vector(MAX_POLYVERT, dtype=float) +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) + +MULTI_CONTACT_COUNT = 4 +mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float) @wp.struct @@ -71,35 +86,55 @@ class Polytope: nhorizon: int +@wp.struct +class SupportPoint: + point: wp.vec3 + cached_index: int + vertex_index: int + + +@wp.func +def _support_margin(geom: Geom, geomtype: int, dir: wp.vec3): + sp = SupportPoint() + sp.cached_index = -1 + sp.vertex_index = -1 + if geomtype == int(GeomType.SPHERE.value): + sp.point = geom.pos + return sp + elif geomtype == int(GeomType.CAPSULE.value): + local_dir = wp.transpose(geom.rot) @ dir + res = wp.vec3() + res[2] = wp.where(local_dir[2] >= 0, geom.size[1], -geom.size[1]) + sp.point = res + return sp + + @wp.func def _support(geom: Geom, geomtype: int, dir: wp.vec3): - cached_index = -1 - vertex_index = -1 + sp = SupportPoint() + sp.cached_index = -1 + sp.vertex_index = -1 local_dir = wp.transpose(geom.rot) @ dir if geomtype == int(GeomType.SPHERE.value): - support_pt = geom.pos + geom.size[0] * dir + sp.point = geom.pos + geom.size[0] * dir elif geomtype == int(GeomType.BOX.value): tmp = wp.sign(local_dir) res = wp.cw_mul(tmp, geom.size) - support_pt = geom.rot @ res + geom.pos - vertex_index = 0 - if tmp[0] > 0: - vertex_index += 1 - if tmp[1] > 0: - vertex_index += 2 - if tmp[2] > 0: - vertex_index += 4 + sp.point = geom.rot @ res + geom.pos + sp.vertex_index = wp.where(tmp[0] > 0, 1, 0) + sp.vertex_index += wp.where(tmp[1] > 0, 2, 0) + sp.vertex_index += wp.where(tmp[2] > 0, 4, 0) elif geomtype == int(GeomType.CAPSULE.value): res = local_dir * geom.size[0] # add cylinder contribution res[2] += wp.sign(local_dir[2]) * geom.size[1] - support_pt = geom.rot @ res + geom.pos + sp.point = geom.rot @ res + geom.pos elif geomtype == int(GeomType.ELLIPSOID.value): res = wp.cw_mul(local_dir, geom.size) res = wp.normalize(res) # transform to ellipsoid res = wp.cw_mul(res, geom.size) - support_pt = geom.rot @ res + geom.pos + sp.point = geom.rot @ res + geom.pos elif geomtype == int(GeomType.CYLINDER.value): res = wp.vec3(0.0, 0.0, 0.0) # set result in XY plane: support on circle @@ -110,23 +145,23 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3): res[1] = local_dir[1] * scl # set result in Z direction res[2] = wp.sign(local_dir[2]) * geom.size[1] - support_pt = geom.rot @ res + geom.pos + sp.point = geom.rot @ res + geom.pos elif geomtype == int(GeomType.MESH.value): max_dist = float(FLOAT_MIN) if geom.graphadr == -1 or geom.vertnum < 10: if geom.index > -1: - cached_index = geom.index + sp.cached_index = geom.index max_dist = wp.dot(geom.vert[geom.index], local_dir) - support_pt = geom.vert[geom.index] + sp.point = geom.vert[geom.index] # exhaustive search over all vertices for i in range(geom.vertnum): vert = geom.vert[geom.vertadr + i] dist = wp.dot(vert, local_dir) if dist > max_dist: max_dist = dist - support_pt = vert - cached_index = geom.vertadr + i - vertex_index = cached_index - geom.vertadr + sp.point = vert + sp.cached_index = geom.vertadr + i + sp.vertex_index = sp.cached_index - geom.vertadr else: numvert = geom.graph[geom.graphadr] vert_edgeadr = geom.graphadr + 2 @@ -137,7 +172,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3): imax = int(0) if geom.index > -1: imax = geom.index - cached_index = geom.index + sp.cached_index = geom.index while True: prev = int(imax) @@ -152,12 +187,12 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3): i += int(1) if imax == prev: break - cached_index = imax + sp.cached_index = imax imax = geom.graph[vert_globalid + imax] - vertex_index = imax - support_pt = geom.vert[geom.vertadr + imax] + sp.vertex_index = imax + sp.point = geom.vert[geom.vertadr + imax] - support_pt = geom.rot @ support_pt + geom.pos + sp.point = geom.rot @ sp.point + geom.pos elif geomtype == int(GeomType.HFIELD.value): max_dist = float(FLOAT_MIN) for i in range(6): @@ -165,10 +200,10 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3): dist = wp.dot(vert, local_dir) if dist > max_dist: max_dist = dist - support_pt = vert - support_pt = geom.rot @ support_pt + geom.pos + sp.point = vert + sp.point = geom.rot @ sp.point + geom.pos - return support_pt, cached_index, vertex_index + return sp @wp.func @@ -193,14 +228,18 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int): @wp.func def _epa_support(pt: Polytope, idx: int, geom1: Geom, geom2: Geom, geom1_type: int, geom2_type: int, dir: wp.vec3): - s1, index1, vertex_index1 = _support(geom1, geom1_type, dir) - s2, index2, vertex_index2 = _support(geom2, geom2_type, -dir) + sp = _support(geom1, geom1_type, dir) + pt.vert1[idx] = sp.point + pt.vert_index1[idx] = sp.vertex_index + index1 = sp.cached_index + + sp = _support(geom2, geom2_type, -dir) + pt.vert2[idx] = sp.point + pt.vert_index2[idx] = sp.vertex_index + index2 = sp.cached_index + + pt.vert[idx] = pt.vert1[idx] - pt.vert2[idx] - pt.vert[idx] = s1 - s2 - pt.vert1[idx] = s1 - pt.vert2[idx] = s2 - pt.vert_index1[idx] = vertex_index1 - pt.vert_index2[idx] = vertex_index2 return index1, index2 @@ -548,6 +587,7 @@ def _gjk( geomtype1: int, geomtype2: int, cutoff: float, + use_margin: bool, ): """Find distance within a tolerance between two geoms.""" cutoff2 = cutoff * cutoff @@ -563,6 +603,9 @@ def _gjk( # set initial guess x_k = x1_0 - x2_0 + use_margin1 = use_margin and (geomtype1 == int(GeomType.SPHERE.value) or geomtype1 == int(GeomType.CAPSULE.value)) + use_margin2 = use_margin and (geomtype2 == int(GeomType.SPHERE.value) or geomtype2 == int(GeomType.CAPSULE.value)) + for k in range(gjk_iterations): xnorm = wp.dot(x_k, x_k) # TODO(kbayes): determine new constant here @@ -570,16 +613,20 @@ def _gjk( break dir_neg = x_k / wp.sqrt(xnorm) + # compute kth support point in geom1 + sp = wp.where(use_margin1, _support_margin(geom1, geomtype1, -dir_neg), _support(geom1, geomtype1, -dir_neg)) + simplex1[n] = sp.point + geom1.index = sp.cached_index + simplex_index1[n] = sp.vertex_index + + # compute kth support point in geom2 + sp = wp.where(use_margin2, _support_margin(geom2, geomtype2, dir_neg), _support(geom2, geomtype2, dir_neg)) + simplex2[n] = sp.point + geom2.index = sp.cached_index + simplex_index2[n] = sp.vertex_index + # compute the kth support point - s1_k, i1, vertex_index1 = _support(geom1, geomtype1, -dir_neg) - s2_k, i2, vertex_index2 = _support(geom2, geomtype2, dir_neg) - geom1.index = i1 - geom2.index = i2 - simplex1[n] = s1_k - simplex2[n] = s2_k - simplex_index1[n] = vertex_index1 - simplex_index2[n] = vertex_index2 - simplex[n] = s1_k - s2_k + simplex[n] = simplex1[n] - simplex2[n] if cutoff == 0.0: if wp.dot(x_k, simplex[n]) > 0: @@ -658,11 +705,7 @@ def _same_side(p0: wp.vec3, p1: wp.vec3, p2: wp.vec3, p3: wp.vec3): n = wp.cross(p1 - p0, p2 - p0) dot1 = wp.dot(n, p3 - p0) dot2 = wp.dot(n, -p0) - if dot1 > 0 and dot2 > 0: - return 1 - if dot1 < 0 and dot2 < 0: - return 1 - return 0 + return (dot1 > 0 and dot2 > 0) or (dot1 < 0 and dot2 < 0) @wp.func @@ -1242,13 +1285,703 @@ def _epa(tolerance2: float, epa_iterations: int, pt: Polytope, geom1: Geom, geom # return from valid face if idx > -1: x1, x2 = _epa_witness(pt, idx) - return -wp.sqrt(pt.face_norm2[idx]), x1, x2 - return 0.0, wp.vec3(), wp.vec3() + return -wp.sqrt(pt.face_norm2[idx]), x1, x2, idx + return 0.0, wp.vec3(), wp.vec3(), -1 + + +# 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)): + v1i = vert_index[face[0]] + v2i = vert_index[face[1]] + v3i = vert_index[face[2]] + + feature_index = wp.vec3i(v1i, v2i, v3i) + feature_vert = wp.mat33() + feature_vert[0] = vert[face[0]] + feature_vert[1] = vert[face[1]] + feature_vert[2] = vert[face[2]] + + if v1i != v2i: + dim = wp.where(v3i == v1i or v3i == v2i, 2, 3) + return dim, feature_index, feature_vert + + feature_index[1] = v3i + feature_vert[1] = vert[face[2]] + + dim = wp.where(v1i != v3i, 2, 1) + return dim, feature_index, feature_vert + + +# 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): + res = wp.vec2i() + for i in range(len1): + for j in range(len2): + if wp.dot(vert1[i], vert2[j]) < -FACE_TOL: + res[0] = i + res[1] = j + return 1, res + return 0, res + + +# 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): + res = wp.vec2i() + for i in range(nface): + for j in range(nedge): + if wp.abs(wp.dot(edge[j], face[i])) < EDGE_TOL: + res[0] = j + res[1] = i + return 1, res + return 0, res + + +# 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): + count = int(0) + res = wp.vec2i() + for i in range(start1, start1 + len1): + for j in range(start2, start2 + len2): + if a1[i] == a2[j]: + res[count] = a1[i] + count += 1 + if count == 2: + return 2, res + return count, res + + +@wp.func +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): + for j in range(start2, start2 + len2): + if a1[i] == a2[j]: + res[count] = a1[i] + count += 1 + if count == 2: + return 2, res + return count, res + + +# compute possible polygon normals of a mesh given up to 3 vertices +@wp.func +def mesh_normals( + # In: + feature_dim: int, + feature_index: wp.vec3i, + mat: wp.mat33, + vertadr: int, + polyadr: int, + polynormal: wp.array(dtype=wp.vec3), + polymapadr: wp.array(dtype=int), + polymapnum: wp.array(dtype=int), + polymap: wp.array(dtype=int), +): + normals = polyverts() + indices = polyindices() + + v1 = feature_index[0] + v2 = feature_index[1] + v3 = feature_index[2] + if feature_dim == 3: + v1_adr = polymapadr[vertadr + v1] + v1_num = polymapnum[vertadr + v1] + + v2_adr = polymapadr[vertadr + v2] + v2_num = polymapnum[vertadr + v2] + + v3_adr = polymapadr[vertadr + v3] + v3_num = polymapnum[vertadr + v3] + + faceset = wp.vec2i() + 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) + if n == 0: + return 0, normals, indices + + # three vertices on mesh define a unique face + normals[0] = mat @ polynormal[polyadr + faceset[0]] + indices[0] = faceset[0] + return 1, normals, indices + + if feature_dim == 2: + v1_adr = polymapadr[vertadr + v1] + v1_num = polymapnum[vertadr + v1] + + v2_adr = polymapadr[vertadr + v2] + 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) + if n == 0: + return 0, normals, indices + for i in range(n): + normals[i] = mat @ polynormal[polyadr + edgeset[i]] + indices[i] = edgeset[i] + return n, normals, indices + + if feature_dim == 1: + v1_adr = polymapadr[vertadr + v1] + v1_num = polymapnum[vertadr + v1] + v1_num = wp.where(v1_num <= MAX_POLYVERT, v1_num, MAX_POLYVERT) + for i in range(v1_num): + index = polymap[v1_adr + i] + normals[i] = mat @ polynormal[polyadr + index] + indices[i] = index + return v1_num, normals, indices + return 0, normals, indices + + +# compute normal directional vectors along possible edges given by up to two vertices +@wp.func +def mesh_edge_normals( + # In: + dim: int, + mat: wp.mat33, + pos: wp.vec3, + vertadr: int, + polyadr: int, + vert: wp.array(dtype=wp.vec3), + polyvertadr: wp.array(dtype=int), + polyvertnum: wp.array(dtype=int), + polyvert: wp.array(dtype=int), + polymapadr: wp.array(dtype=int), + polymapnum: wp.array(dtype=int), + polymap: wp.array(dtype=int), + v1: wp.vec3, + v2: wp.vec3, + v1i: int, +): + normals = polyverts() + endverts = polyverts() + + # only one edge + if dim == 2: + endverts[0] = v2 + normals[0] = wp.normalize(v2 - v1) + return 1, normals, endverts + + if dim == 1: + v1_adr = polymapadr[vertadr + v1i] + v1_num = polymapnum[vertadr + v1i] + v1_num = wp.where(v1_num <= MAX_POLYVERT, v1_num, MAX_POLYVERT) + + # loop through all faces with vertex v1 + for i in range(v1_num): + idx = polymap[v1_adr + i] + adr = polyvertadr[polyadr + idx] + nvert = polyvertnum[polyadr + idx] + # find previous vertex in polygon to form edge + for j in range(nvert): + if polyvert[adr + j] == v1i: + k = wp.where(j == 0, nvert - 1, j - 1) + endverts[i] = mat @ vert[vertadr + polyvert[adr + k]] + pos + normals[i] = wp.normalize(endverts[i] - v1) + return v1_num, normals, endverts + return 0, normals, endverts + + +# try recovering box normal from collision normal +@wp.func +def box_normals2(mat: wp.mat33, n: wp.vec3): + normals = polyverts() + indices = polyindices() + + # list of box face normals + face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0) + + # get local coordinates of the normal + + local_n = wp.normalize( + wp.vec3( + mat[0][0] * n[0] + mat[1][0] * n[1] + mat[2][0] * n[2], + mat[0][1] * n[0] + mat[1][1] * n[1] + mat[2][1] * n[2], + mat[0][2] * n[0] + mat[1][2] * n[1] + mat[2][2] * n[2], + ) + ) + + # determine if there is a side close to the normal + for i in range(6): + if wp.dot(local_n, face_normals[i]) > FACE_TOL: + normals[0] = mat @ face_normals[i] + indices[0] = i + return 1, normals, indices + + return 0, normals, indices + + +# 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): + normals = polyverts() + indices = polyindices() + + v1 = feature_index[0] + v2 = feature_index[1] + v3 = feature_index[2] + + if feature_dim == 3: + c = 0 + x = float((v1 & 1) and (v2 & 1) and (v3 & 1)) - float(not (v1 & 1) and not (v2 & 1) and not (v3 & 1)) + y = float((v1 & 2) and (v2 & 2) and (v3 & 2)) - float(not (v1 & 2) and not (v2 & 2) and not (v3 & 2)) + z = float((v1 & 4) and (v2 & 4) and (v3 & 4)) - float(not (v1 & 4) and not (v2 & 4) and not (v3 & 4)) + normals[0] = mat @ wp.vec3(x, y, z) + sgn = x + y + z + if x != 0.0: + indices[c] = 0 + c += 1 + if y != 0.0: + indices[c] = 2 + c += 1 + if z != 0.0: + indices[c] = 4 + c += 1 + if sgn == -1.0: + indices[0] = indices[0] + 1 + if c == 1: + return 1, normals, indices + 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)) + y = float((v1 & 2) and (v2 & 2)) - float(not (v1 & 2) and not (v2 & 2)) + z = float((v1 & 4) and (v2 & 4)) - float(not (v1 & 4) and not (v2 & 4)) + if x != 0.0: + normals[c] = mat @ wp.vec3(float(x), 0.0, 0.0) + indices[c] = wp.where(x > 0.0, 0, 1) + c += 1 + if y != 0.0: + normals[c] = mat @ wp.vec3(0.0, y, 0.0) + indices[c] = wp.where(y > 0.0, 2, 3) + c += 1 + if z != 0.0: + normals[c] = mat @ wp.vec3(0.0, 0.0, z) + indices[c] = wp.where(z > 0.0, 4, 5) + c += 1 + if c == 2: + return 2, normals, indices + return box_normals2(mat, dir) + + if feature_dim == 1: + x = wp.where(v1 & 1, 1.0, -1.0) + y = wp.where(v1 & 2, 1.0, -1.0) + z = wp.where(v1 & 4, 1.0, -1.0) + normals[0] = mat @ wp.vec3(x, 0.0, 0.0) + normals[1] = mat @ wp.vec3(0.0, y, 0.0) + normals[2] = mat @ wp.vec3(0.0, 0.0, z) + indices[0] = wp.where(x > 0.0, 0, 1) + indices[1] = wp.where(y > 0.0, 2, 3) + indices[2] = wp.where(z > 0.0, 4, 5) + return 3, normals, indices + return 0, normals, indices + + +# 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): + normals = polyverts() + endverts = polyverts() + + if dim == 2: + endverts[0] = v2 + normals[0] = wp.normalize(v2 - v1) + return 1, normals, endverts + + # return 3 adjacent vertices + if dim == 1: + x = wp.where(v1i & 1, size[0], -size[0]) + y = wp.where(v1i & 2, size[1], -size[1]) + z = wp.where(v1i & 4, size[2], -size[2]) + + endverts[0] = mat @ wp.vec3(-x, y, z) + pos + normals[0] = wp.normalize(endverts[0] - v1) + + endverts[1] = mat @ wp.vec3(x, -y, z) + pos + normals[1] = wp.normalize(endverts[1] - v1) + + endverts[2] = mat @ wp.vec3(x, y, -z) + pos + normals[2] = wp.normalize(endverts[2] - v1) + return 3, normals, endverts + return 0, normals, endverts + + +# recover face of a box from its index +@wp.func +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 + if idx == 0: # right + res[0] = mat @ wp.vec(size[0], size[1], size[2]) + pos + res[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos + res[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos + res[3] = mat @ wp.vec(size[0], -size[1], size[2]) + pos + return 4, res + if idx == 1: # left + res[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos + res[1] = mat @ wp.vec(-size[0], size[1], size[2]) + pos + res[2] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos + res[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + return 4, res + if idx == 2: # top + res[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos + res[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos + res[2] = mat @ wp.vec(size[0], size[1], size[2]) + pos + res[3] = mat @ wp.vec(-size[0], size[1], size[2]) + pos + return 4, res + if idx == 3: # bottom + res[0] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos + res[1] = mat @ wp.vec(size[0], -size[1], size[2]) + pos + res[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos + res[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + return 4, res + if idx == 4: # front + res[0] = mat @ wp.vec(-size[0], size[1], size[2]) + pos + res[1] = mat @ wp.vec(size[0], size[1], size[2]) + pos + res[2] = mat @ wp.vec(size[0], -size[1], size[2]) + pos + res[3] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos + return 4, res + if idx == 5: # back + res[0] = mat @ wp.vec(size[0], size[1], -size[2]) + pos + res[1] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos + res[2] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + res[3] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos + return 4, res + return 0, res + + +# recover mesh polygon from its index, return number of edges +@wp.func +def mesh_face( + # In: + mat: wp.mat33, + pos: wp.vec3, + vertadr: int, + polyadr: int, + vert: wp.array(dtype=wp.vec3), + polyvertadr: wp.array(dtype=int), + polyvertnum: wp.array(dtype=int), + polyvert: wp.array(dtype=int), + idx: int, +): + res = polyverts() + + adr = polyvertadr[polyadr + idx] + j = int(0) + nvert = polyvertnum[polyadr + idx] + nvert = wp.where(nvert <= MAX_POLYVERT, nvert, MAX_POLYVERT) + for i in range(nvert - 1, -1, -1): + v = vert[vertadr + polyvert[adr + i]] + res[j] = mat @ v + pos + j += 1 + return nvert, res + + +@wp.func +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): + return wp.dot(p - a, n) > -MJ_MINVAL + + +@wp.func +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) + if temp == 0.0: + return FLOAT_MAX, res # parallel; no intersection + t = (pd - wp.dot(pn, a)) / temp + if t >= 0.0 and t <= 1.0: + res[0] = a[0] + t * ab[0] + res[1] = a[1] + t * ab[1] + res[2] = a[2] + t * ab[2] + return t, res + + +# 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): + witness1 = mat3c() + witness2 = mat3c() + + # clipping face needs to be at least a triangle + if nface1 < 3: + return 0, witness1, witness2 + + # compute plane normal and distance to plane for each vertex + pn = polyverts() + pd = polyvec() + for i in range(nface1): + 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) + pd[nface1 - 1] = pdi + pn[nface1 - 1] = pni + + # reserve 2 * max_sides as max sides for a clipped polygon + polygon1 = polyclip() + polygon2 = polyclip() + npolygon = nface2 + nclipped = int(0) + + polygon = polygon1 + clipped = polygon2 + + for i in range(nface2): + polygon[i] = face2[i] + + # clip the polygon by one edge e at a time + for e in range(nface1): + for i in range(npolygon): + # get edge PQ of the polygon + P = polygon[i] + 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) + + # PQ entirely outside the clipping edge, skip + if not inside1 and not inside2: + continue + + # edge PQ is inside the clipping edge, add Q + if inside1 and inside2: + clipped[nclipped] = Q + nclipped += 1 + continue + + # add new vertex to clipped polygon where PQ intersects the clipping edge + t, res = plane_intersect(pn[e], pd[e], P, Q) + if t < 0.0 or t > 1.0: + clipped[nclipped] = res + nclipped += 1 + + # add Q as PQ is now back inside the clipping edge + if inside2: + clipped[nclipped] = Q + nclipped += 1 + + # swap clipped and polygon + tmp = polygon + polygon = clipped + clipped = tmp + npolygon = nclipped + nclipped = 0 + + if npolygon < 1: + return 0, witness1, witness2 + + # no pruning needed + for i in range(npolygon): + witness2[i] = polygon[i] + witness1[i] = witness2[i] + dir + return npolygon, witness2, witness1 + + +# recover multiple contacts from EPA polytope +@wp.func +def multicontact( + pt: Polytope, face: wp.vec3i, x1: wp.vec3, x2: wp.vec3, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int +): + witness1 = mat3c() + witness2 = mat3c() + witness1[0] = x1 + witness2[0] = x2 + + face1 = polyverts() + face2 = polyverts() + endverts = polyverts() + + if geomtype1 == int(GeomType.MESH.value): + vert = geom1.vert + polynormal = geom1.mesh_polynormal + polyvertadr = geom1.mesh_polyvertadr + polyvertnum = geom1.mesh_polyvertnum + polyvert = geom1.mesh_polyvert + polymapadr = geom1.mesh_polymapadr + polymapnum = geom1.mesh_polymapnum + polymap = geom1.mesh_polymap + elif geomtype2 == int(GeomType.MESH.value): + vert = geom2.vert + polynormal = geom2.mesh_polynormal + polyvertadr = geom2.mesh_polyvertadr + polyvertnum = geom2.mesh_polyvertnum + polyvert = geom2.mesh_polyvert + polymapadr = geom2.mesh_polymapadr + polymapnum = geom2.mesh_polymapnum + 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) + + 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) + elif geomtype1 == int(GeomType.MESH.value): + 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) + elif geomtype2 == int(GeomType.MESH.value): + 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) + 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( + 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( + nface1, + geom1.rot, + geom1.pos, + geom1.vertadr, + geom1.mesh_polyadr, + geom1.vert, + polyvertadr, + polyvertnum, + polyvert, + polymapadr, + polymapnum, + polymap, + feature_vertex1[0], + feature_vertex1[1], + feature_index1[0], + ) + nres, res = aligned_face_edge(n1, nnorms1, n2, nnorms2) + if not nres: + return 1, witness1, witness2 + is_edge_contact_geom1 = 1 + + # check if face-edge collision + elif nface2 < 3: + nnorms2 = 0 + if geomtype2 == int(GeomType.BOX.value): + 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( + nface2, + geom2.rot, + geom2.pos, + geom2.vertadr, + geom2.mesh_polyadr, + geom2.vert, + polyvertadr, + polyvertnum, + polyvert, + polymapadr, + polymapnum, + polymap, + feature_vertex2[0], + feature_vertex2[1], + feature_index2[0], + ) + nres, res = aligned_face_edge(n2, nnorms2, n1, nnorms1) + if not nres: + return 1, witness1, witness2 + is_edge_contact_geom2 = 1 + else: + # no multi-contact + return 1, witness1, witness2 + + i = res[0] + j = res[1] + + # recover geom1 matching edge or face + if is_edge_contact_geom1: + face1[0] = pt.vert1[face[0]] + face1[1] = endverts[i] + nface1 = 2 + 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) + elif geomtype1 == int(GeomType.MESH.value): + nface1, face1 = mesh_face( + geom1.rot, geom1.pos, geom1.vertadr, geom1.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, ind + ) + + # recover geom2 matching edge or face + if is_edge_contact_geom2: + face2[0] = pt.vert2[face[0]] + face2[1] = endverts[i] + nface2 = 2 + else: + if geomtype2 == int(GeomType.BOX.value): + nface2, face2 = box_face(geom2.rot, geom2.pos, geom2.size, idx2[j]) + elif geomtype2 == int(GeomType.MESH.value): + nface2, face2 = mesh_face( + geom2.rot, geom2.pos, geom2.vertadr, geom2.mesh_polyadr, vert, polyvertadr, polyvertnum, polyvert, idx2[j] + ) + + # TODO(kbayes): this approximates the contact direction, by scaling the face normal by the + # single contact direction's magnitude. This is effective, but polygonClip should compute + # this for each contact point. + approx_dir = wp.vec3() + + # 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) + + # 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) + + # face-face collision + approx_dir = wp.norm_l2(dir) * n2[j] + 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): + n = wp.normalize(x2 - x1) + if margin1 > 0.0: + x1 += margin1 * n + + if margin2 > 0.0: + x2 -= margin2 * n + dist -= margin1 + margin2 + return dist, x1, x2 @wp.func def ccd( # In: + multiccd: bool, tolerance: float, cutoff: float, gjk_iterations: int, @@ -1271,12 +2004,44 @@ def ccd( face_map: wp.array(dtype=int), horizon: wp.array(dtype=int), ): + witness1 = mat3c() + witness2 = mat3c() """General convex collision detection via GJK/EPA.""" - result = _gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff) + margin1 = 0.0 + margin2 = 0.0 + + if geomtype1 == int(GeomType.SPHERE.value) or geomtype1 == int(GeomType.CAPSULE.value): + margin1 = geom1.size[0] + + if geomtype2 == int(GeomType.SPHERE.value) or geomtype2 == int(GeomType.CAPSULE.value): + margin2 = geom2.size[0] + + # special handling for sphere and capsule (shrink to point and line respectively) + if margin1 + margin2 > 0.0: + cutoff += margin1 + margin2 + result = _gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, True) + + # shallow penetration, inflate contact + if result.dist > tolerance: + if result.dist == FLOAT_MAX: + 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) + witness1[0] = x1 + witness2[0] = x2 + return dist, 1, witness1, witness2 + + # deep penetration, reset initial conditions and rerun GJK + EPA + cutoff -= margin1 + margin2 + + result = _gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, False) # no penetration depth to recover if result.dist > tolerance or result.dim < 2: - return result.dist, result.x1, result.x2 + witness1[0] = result.x1 + witness2[0] = result.x2 + return result.dist, 1, witness1, witness2 pt = Polytope() pt.nface = 0 @@ -1356,6 +2121,19 @@ def ccd( # origin on boundary (objects are not considered penetrating) if pt.status: - return result.dist, result.x1, result.x2 + witness1[0] = result.x1 + witness2[0] = result.x2 + return result.dist, 1, witness1, witness2 - return _epa(tolerance * tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2) + dist, x1, x2, idx = _epa(tolerance * tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2) + if ( + multiccd + and (geomtype1 == int(GeomType.BOX.value) or geomtype1 == int(GeomType.MESH.value)) + and (geomtype2 == int(GeomType.BOX.value) or geomtype2 == int(GeomType.MESH.value)) + ): + 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 + witness2[0] = x2 + return dist, 1, witness1, witness2 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 8ccb259a..a490b862 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 @@ -96,9 +96,11 @@ def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int): ( dist, + count, x1, x2, ) = ccd( + False, 1e-6, 1.0e30, iterations, @@ -123,8 +125,8 @@ def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int): ) dist_out[0] = dist - pos_out[0] = x1 - pos_out[1] = x2 + pos_out[0] = x1[0] + pos_out[1] = x2[0] vert = wp.array(shape=(iterations,), dtype=wp.vec3) vert1 = wp.array(shape=(iterations,), dtype=wp.vec3) @@ -253,8 +255,7 @@ class GJKTest(absltest.TestCase): """ ) - # TODO(kbayes): use margin trick instead of EPA for penetration recovery - dist, _, _ = _geom_dist(m, d, 0, 1, 500) + dist, _, _ = _geom_dist(m, d, 0, 1, 0) self.assertAlmostEqual(-2, dist) def test_box_box_contact(self): diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py index 88457e91..d00ccdb1 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py @@ -65,11 +65,33 @@ def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp. return aabb +@wp.func +def radial_field(a: wp.vec3, x: wp.vec3, size: wp.vec3) -> wp.vec3: + field = wp.cw_div(-size, a) + field = wp.normalize(field) + field[0] *= wp.sign(x[0]) + field[1] *= wp.sign(x[1]) + field[2] *= wp.sign(x[2]) + return field + + @wp.func def sphere(p: wp.vec3, size: wp.vec3) -> float: return wp.length(p) - size[0] +@wp.func +def box(p: wp.vec3, size: wp.vec3) -> float: + a = wp.abs(p) - size + if a[0] >= 0 or a[1] >= 0 or a[2] >= 0: + z = wp.vec3(0.0, 0.0, 0.0) + b = wp.max(a, z) + return wp.norm_l2(b) + wp.min(wp.max(a), 0.0) + b = radial_field(a, p, size) + t = -wp.cw_div(a, wp.abs(b)) + return -wp.min(t) * wp.norm_l2(b) + + @wp.func def ellipsoid(p: wp.vec3, size: wp.vec3) -> float: scaled_p = wp.vec3(p[0] / size[0], p[1] / size[1], p[2] / size[2]) @@ -91,6 +113,24 @@ def grad_sphere(p: wp.vec3) -> wp.vec3: wp.vec3(0.0) +@wp.func +def grad_box(p: wp.vec3, size: wp.vec3) -> wp.vec3: + a = wp.abs(p) - size + if wp.max(a) < 0: + return radial_field(a, p, size) + z = wp.vec3(0.0, 0.0, 0.0) + b = wp.max(a, z) + c = wp.norm_l2(b) + g = wp.cw_mul(wp.div(b, c), wp.cw_div(p, wp.abs(p))) + if a[0] <= 0: + g[0] = 0.0 + if a[1] <= 0: + g[1] = 0.0 + if a[2] <= 0: + g[2] = 0.0 + return g + + @wp.func def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3: a = wp.vec3(p[0] / size[0], p[1] / size[1], p[2] / size[2]) @@ -128,8 +168,12 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3: @wp.func def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float: - if type == int(GeomType.SPHERE.value): + if type == int(GeomType.PLANE.value): + return p[2] + elif type == int(GeomType.SPHERE.value): return sphere(p, attr) + elif type == int(GeomType.BOX.value): + return box(p, attr) elif type == int(GeomType.ELLIPSOID.value): return ellipsoid(p, attr) elif type == int(GeomType.SDF.value): @@ -140,8 +184,13 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float: @wp.func def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3: - if type == int(GeomType.SPHERE.value): + if type == int(GeomType.PLANE.value): + grad = wp.vec3(0.0, 0.0, 1.0) + return grad + elif type == int(GeomType.SPHERE.value): return grad_sphere(p) + elif type == int(GeomType.BOX.value): + return grad_box(p, attr) elif type == int(GeomType.ELLIPSOID.value): return grad_ellipsoid(p, attr) elif type == int(GeomType.SDF.value): @@ -444,41 +493,40 @@ def _sdf_narrowphase( g1_plugin = geom_plugin_index[g1] g2_plugin = geom_plugin_index[g2] - g2_to_g1_rot = wp.transpose(geom2.rot) * geom1.rot - g2_to_g1_pos = wp.transpose(geom2.rot) * (geom1.pos - geom2.pos) + g1_to_g2_rot = wp.transpose(geom1.rot) * geom2.rot + g1_to_g2_pos = wp.transpose(geom1.rot) * (geom2.pos - geom1.pos) aabb_pos = geom_aabb[g1, 0] aabb_size = geom_aabb[g1, 1] - aabb1 = transform_aabb(aabb_pos, aabb_size, g2_to_g1_pos, g2_to_g1_rot) - + identity = wp.identity(3, dtype=float) + aabb1 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), identity) aabb_pos = geom_aabb[g2, 0] aabb_size = geom_aabb[g2, 1] - aabb2 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), wp.mat33(1.0)) + aabb2 = transform_aabb(aabb_pos, aabb_size, g1_to_g2_pos, g1_to_g2_rot) aabb_intersection = AABB() aabb_intersection.min = wp.max(aabb1.min, aabb2.min) aabb_intersection.max = wp.min(aabb1.max, aabb2.max) - geom_pos2 = geom_pos[worldid, g2] - quat2 = geom_quat[worldid, g2] - geom_mat2 = math.quat_to_mat(quat2) - rot2 = math.mul(geom2.rot, math.transpose(geom_mat2)) - pos2 = wp.sub(geom2.pos, math.mul(rot2, geom_pos2)) + pos2 = geom2.pos + rot2 = geom2.rot + pos1 = geom1.pos + rot1 = geom1.rot if type1 == int(GeomType.SDF.value): - geom_pos1 = geom_pos[worldid, g1] - quat1 = geom_quat[worldid, g1] - geom_mat1 = math.quat_to_mat(quat1) - rot1 = math.mul(geom1.rot, math.transpose(geom_mat1)) - pos1 = wp.sub(geom1.pos, math.mul(rot1, geom_pos1)) attr1 = plugin_attr[g1_plugin] g1_plugin_id = plugin[g1_plugin] else: - pos1 = geom1.pos - rot1 = geom1.rot attr1 = geom1.size g1_plugin_id = -1 + if g2_plugin != -1: + attr2 = plugin_attr[g2_plugin] + g2_plugin_id = plugin[g2_plugin] + else: + attr2 = geom2.size + g2_plugin_id = -1 + for i in range(sdf_initpoints): x_g2 = wp.vec3( aabb_intersection.min[0] + (aabb_intersection.max[0] - aabb_intersection.min[0]) * halton(i, 2), @@ -486,11 +534,11 @@ def _sdf_narrowphase( aabb_intersection.min[2] + (aabb_intersection.max[2] - aabb_intersection.min[2]) * halton(i, 5), ) - x = geom2.rot * x_g2 + geom2.pos + x = geom1.rot * x_g2 + geom1.pos x0_initial = wp.transpose(rot2) * (x - pos2) dist, pos, n = gradient_descent( - type1, x0_initial, attr1, plugin_attr[g2_plugin], pos1, rot1, pos2, rot2, g1_plugin_id, plugin[g2_plugin], sdf_iterations + type1, x0_initial, attr1, attr2, pos1, rot1, pos2, rot2, g1_plugin_id, g2_plugin_id, sdf_iterations ) write_contact( diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py index 28cffc8c..036d2052 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py @@ -65,33 +65,32 @@ class ConstraintTest(parameterized.TestCase): xml = f""" - - + - - + + + + + + + + """ - _, mjd, m, d = test_util.fixture(xml=xml, cone=cone) - - for arr in ( - d.efc.D, - d.efc.aref, - d.efc.pos, - d.efc.margin, - ): - arr.zero_() + _, mjd, m, d = test_util.fixture(xml=xml, cone=cone, keyframe=0) # fill with nan to check whether we are not reading uninitialized values - d.efc.J.fill_(wp.nan) + for arr in (d.efc.J, d.efc.D, d.efc.aref, d.efc.pos, d.efc.margin): + arr.fill_(wp.nan) mjwarp.make_constraint(m, d) + _assert_eq(d.ncon.numpy()[0], mjd.ncon, "ncon") _assert_eq(d.efc.J.numpy()[0, : mjd.nefc, :].reshape(-1), mjd.efc_J, "efc_J") _assert_eq(d.efc.D.numpy()[0, : mjd.nefc], mjd.efc_D, "efc_D") _assert_eq(d.efc.aref.numpy()[0, : mjd.nefc], mjd.efc_aref, "efc_aref") @@ -105,21 +104,12 @@ class ConstraintTest(parameterized.TestCase): def test_constraints(self, cone): """Test constraints.""" for key in range(3): - mjm, mjd, m, d = test_util.fixture("constraints.xml", sparse=False, cone=cone, keyframe=key) + _, mjd, m, d = test_util.fixture("constraints.xml", sparse=False, cone=cone, keyframe=key) - for arr in ( - d.efc.D, - d.efc.aref, - d.efc.pos, - d.efc.margin, - d.ne, - d.nefc, - d.nf, - d.nl, - ): - arr.zero_() - - d.efc.J.fill_(wp.nan) + for arr in (d.ne, d.nefc, d.nf, d.nl, d.efc.type): + arr.fill_(-1) + for arr in (d.efc.J, d.efc.D, d.efc.aref, d.efc.pos, d.efc.margin): + arr.fill_(wp.nan) mjwarp.make_constraint(m, d) @@ -140,8 +130,10 @@ class ConstraintTest(parameterized.TestCase): for keyframe in range(-1, 1): _, mjd, m, d = test_util.fixture("tendon/tendon_limit.xml", sparse=False, keyframe=keyframe) - for arr in (d.nefc, d.nl, d.efc.J, d.efc.D, d.efc.aref, d.efc.pos, d.efc.margin): - arr.zero_() + for arr in (d.nefc, d.nl, d.efc.type): + arr.fill_(-1) + for arr in (d.efc.J, d.efc.D, d.efc.aref, d.efc.pos, d.efc.margin): + arr.fill_(wp.nan) mjwarp.make_constraint(m, d) @@ -197,9 +189,15 @@ class ConstraintTest(parameterized.TestCase): - """ + """, + keyframe=0, ) + for arr in (d.nefc, d.ne, d.efc.type): + arr.fill_(-1) + for arr in (d.efc.J, d.efc.D, d.efc.vel, d.efc.aref, d.efc.pos, d.efc.margin): + arr.fill_(wp.nan) + mjwarp.make_constraint(m, d) _assert_eq(d.nefc.numpy()[0], mjd.nefc, "nefc") 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 a27c38a1..ad07b5c4 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -1017,9 +1017,9 @@ def forward(m: Model, d: Data): fwd_actuation(m, d) fwd_acceleration(m, d, factorize=True) - sensor.sensor_acc(m, d) solver.solve(m, d) + sensor.sensor_acc(m, d) @event_scope 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 f2f2faa5..0c9dbb25 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -27,6 +27,9 @@ from mujoco.mjx.third_party.mujoco_warp._src import types # number of max iterations to run GJK/EPA MJ_CCD_ITERATIONS = 12 +# max number of worlds supported +MAX_WORLDS = 2**24 + def _hfield_geom_pair(mjm: mujoco.MjModel) -> Tuple[int, np.array]: geom1, geom2 = np.triu_indices(mjm.ngeom, k=1) @@ -115,6 +118,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: if mjm.opt.noslip_iterations > 0: raise NotImplementedError(f"noslip solver not implemented.") + # contact sensor + is_contact_sensor = mjm.sensor_type == types.SensorType.CONTACT + if is_contact_sensor.any(): + # matching + if ( + (mjm.sensor_objtype[is_contact_sensor] != types.ObjType.GEOM) + | (mjm.sensor_reftype[is_contact_sensor] != types.ObjType.GEOM) + ).any(): + 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.") + # TODO(team): remove after _update_gradient for Newton uses tile operations for islands nv_max = 60 if mjm.nv > nv_max and mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_DENSE: @@ -366,10 +383,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: array._is_batched = True if not expand_dim: array.strides = (0,) + array.strides[1:] + array.shape = (MAX_WORLDS,) + array.shape[1:] return array array.strides = (0,) + array.strides array.ndim += 1 - array.shape = (1,) + array.shape + array.shape = (MAX_WORLDS,) + array.shape return array # rangefinder @@ -415,6 +433,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: nwrap=mjm.nwrap, nsensor=mjm.nsensor, nsensordata=mjm.nsensordata, + nsensortaxel=sum(mjm.mesh_vertnum[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_TACTILE]]), nmeshvert=mjm.nmeshvert, nmeshface=mjm.nmeshface, nmeshgraph=mjm.nmeshgraph, @@ -595,10 +614,13 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: mesh_vertadr=wp.array(mjm.mesh_vertadr, dtype=int), mesh_vertnum=wp.array(mjm.mesh_vertnum, dtype=int), mesh_vert=wp.array(mjm.mesh_vert, dtype=wp.vec3), + mesh_normaladr=wp.array(mjm.mesh_normaladr, dtype=int), + mesh_normal=wp.array(mjm.mesh_normal, dtype=wp.vec3), mesh_faceadr=wp.array(mjm.mesh_faceadr, dtype=int), mesh_face=wp.array(mjm.mesh_face, dtype=wp.vec3i), mesh_graphadr=wp.array(mjm.mesh_graphadr, dtype=int), mesh_graph=wp.array(mjm.mesh_graph, dtype=int), + mesh_quat=wp.array(mjm.mesh_quat, dtype=wp.quat), mesh_polynum=wp.array(mjm.mesh_polynum, dtype=int), mesh_polyadr=wp.array(mjm.mesh_polyadr, dtype=int), mesh_polynormal=wp.array(mjm.mesh_polynormal, dtype=wp.vec3), @@ -709,6 +731,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: sensor_objid=wp.array(mjm.sensor_objid, dtype=int), sensor_reftype=wp.array(mjm.sensor_reftype, dtype=int), sensor_refid=wp.array(mjm.sensor_refid, dtype=int), + sensor_intprm=wp.array(mjm.sensor_intprm, dtype=int), sensor_dim=wp.array(mjm.sensor_dim, dtype=int), sensor_adr=wp.array(mjm.sensor_adr, dtype=int), sensor_cutoff=wp.array(mjm.sensor_cutoff, dtype=float), @@ -776,6 +799,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: mjm.sensor_type, [mujoco.mjtSensor.mjSENS_SUBTREELINVEL, mujoco.mjtSensor.mjSENS_SUBTREEANGMOM], ).any(), + sensor_contact_adr=wp.array(np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT)[0], dtype=int), sensor_rne_postconstraint=np.isin( mjm.sensor_type, [ @@ -798,6 +822,24 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: block_dim=types.BlockDim(), geom_pair_type_count=tuple(geom_type_pair_count), has_sdf_geom=bool(np.any(mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF)), + taxel_vertadr=wp.array( + [ + j + mjm.mesh_vertadr[mjm.sensor_objid[i]] + for i in range(mjm.nsensor) + if mjm.sensor_type[i] == mujoco.mjtSensor.mjSENS_TACTILE + for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]]) + ], + dtype=int, + ), + taxel_sensorid=wp.array( + [ + i + for i in range(mjm.nsensor) + if mjm.sensor_type[i] == mujoco.mjtSensor.mjSENS_TACTILE + for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]]) + ], + dtype=int, + ), ) return m @@ -811,11 +853,12 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in mjm (mujoco.MjModel): The model containing kinematic and dynamic information (host). nworld (int, optional): Number of worlds. Defaults to 1. nconmax (int, optional): Maximum number of contacts for all worlds. Defaults to -1. - njmax (int, optional): Maximum number of constraints for all worlds. Defaults to -1. + njmax (int, optional): Maximum number of constraints per world. Defaults to -1. Returns: Data: The data object containing the current state and output arrays (device). """ + # TODO(team): move to Model? if nconmax == -1: # TODO(team): heuristic for nconmax @@ -823,6 +866,16 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in if njmax == -1: # TODO(team): heuristic for njmax njmax = 20 * 6 + + if nworld < 1 or nworld > MAX_WORLDS: + raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}") + + if nconmax < 1: + raise ValueError("nconmax must be >= 1") + + if njmax < 1: + raise ValueError("njmax must be >= 1") + condim = np.concatenate((mjm.geom_condim, mjm.pair_dim)) condim_max = np.max(condim) if len(condim) > 0 else 0 @@ -833,6 +886,7 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in qM = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float) qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float) + nsensorcontact = np.sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT) nrangefinder = sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER) return types.Data( @@ -841,6 +895,7 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in njmax=njmax, solver_niter=wp.zeros(nworld, dtype=int), ncon=wp.zeros(1, dtype=int), + ncon_world=wp.zeros(nworld, dtype=int), ncon_hfield=wp.zeros((nworld, _hfield_geom_pair(mjm)[0]), dtype=int), # warp only ne=wp.zeros(nworld, dtype=int), ne_connect=wp.zeros(nworld, dtype=int), # warp only @@ -965,8 +1020,6 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in prev_grad=wp.zeros((nworld, mjm.nv), dtype=float), prev_Mgrad=wp.zeros((nworld, mjm.nv), dtype=float), beta=wp.zeros((nworld,), dtype=float), - beta_num=wp.zeros((nworld,), dtype=float), - beta_den=wp.zeros((nworld,), dtype=float), done=wp.zeros((nworld,), dtype=bool), # linesearch ls_done=wp.zeros((nworld,), dtype=bool), @@ -1051,6 +1104,10 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in sensor_rangefinder_vec=wp.zeros((nworld, nrangefinder), dtype=wp.vec3), sensor_rangefinder_dist=wp.zeros((nworld, nrangefinder), dtype=float), sensor_rangefinder_geomid=wp.zeros((nworld, nrangefinder), dtype=int), + sensor_contact_nmatch=wp.zeros((nworld, nsensorcontact), dtype=int), + sensor_contact_matchid=wp.zeros((nworld, nsensorcontact, types.MJ_MAXCONPAIR), dtype=int), + sensor_contact_criteria=wp.zeros((nworld, nsensorcontact, types.MJ_MAXCONPAIR), dtype=float), + sensor_contact_direction=wp.zeros((nworld, nsensorcontact, types.MJ_MAXCONPAIR), dtype=float), # ray ray_bodyexclude=wp.zeros(1, dtype=int), ray_dist=wp.zeros((nworld, 1), dtype=float), @@ -1078,7 +1135,7 @@ def put_data( mjd (mujoco.MjData): The data object containing current state and output arrays (host). nworld (int, optional): The number of worlds. Defaults to 1. nconmax (int, optional): The maximum number of contacts for all worlds. Defaults to -1. - njmax (int, optional): The maximum number of constraints for all worlds. Defaults to -1. + njmax (int, optional): The maximum number of constraints per world. Defaults to -1. Returns: Data: The data object containing the current state and output arrays (device). @@ -1089,12 +1146,12 @@ def put_data( nworld = nworld or 1 # TODO(team): better heuristic for nconmax - nconmax = nconmax or max(512, mjd.ncon * nworld) + nconmax = nconmax or max(512, 4 * mjd.ncon * nworld) # TODO(team): better heuristic for njmax - njmax = njmax or max(5, mjd.nefc) + njmax = njmax or max(5, 4 * mjd.nefc) - if nworld < 1: - raise ValueError("nworld must be >= 1") + if nworld < 1 or nworld > MAX_WORLDS: + raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}") if nconmax < 1: raise ValueError("nconmax must be >= 1") @@ -1192,6 +1249,7 @@ def put_data( efc_force_fill[:, :nefc] = np.tile(mjd.efc_force, (nworld, 1)) efc_margin_fill[:, :nefc] = np.tile(mjd.efc_margin, (nworld, 1)) + nsensorcontact = np.sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT) nrangefinder = sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER) # some helper functions to simplify the data field definitions below @@ -1225,6 +1283,7 @@ def put_data( njmax=njmax, solver_niter=tile(mjd.solver_niter[0]), ncon=arr([mjd.ncon * nworld]), + ncon_world=wp.zeros(nworld, dtype=int), ncon_hfield=wp.zeros((nworld, _hfield_geom_pair(mjm)[0]), dtype=int), # warp only ne=wp.full(shape=(nworld), value=mjd.ne), ne_connect=wp.full(shape=(nworld), value=ne_connect), @@ -1346,8 +1405,6 @@ def put_data( prev_grad=wp.empty(shape=(nworld, mjm.nv), dtype=float), prev_Mgrad=wp.empty(shape=(nworld, mjm.nv), dtype=float), beta=wp.empty(shape=(nworld,), dtype=float), - beta_num=wp.empty(shape=(nworld,), dtype=float), - beta_den=wp.empty(shape=(nworld,), dtype=float), done=wp.empty(shape=(nworld,), dtype=bool), ls_done=wp.zeros(shape=(nworld,), dtype=bool), p0=wp.empty(shape=(nworld,), dtype=wp.vec3), @@ -1429,6 +1486,10 @@ def put_data( sensor_rangefinder_vec=wp.zeros((nworld, nrangefinder), dtype=wp.vec3), sensor_rangefinder_dist=wp.zeros((nworld, nrangefinder), dtype=float), sensor_rangefinder_geomid=wp.zeros((nworld, nrangefinder), dtype=int), + sensor_contact_nmatch=wp.zeros((nworld, nsensorcontact), dtype=int), + sensor_contact_matchid=wp.zeros((nworld, nsensorcontact, types.MJ_MAXCONPAIR), dtype=int), + sensor_contact_criteria=wp.zeros((nworld, nsensorcontact, types.MJ_MAXCONPAIR), dtype=float), + sensor_contact_direction=wp.zeros((nworld, nsensorcontact, types.MJ_MAXCONPAIR), dtype=float), # ray ray_bodyexclude=wp.zeros(1, dtype=int), ray_dist=wp.zeros((nworld, 1), dtype=float), 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 33503268..8599d025 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 @@ -23,10 +23,12 @@ import mujoco import numpy as np import warp as wp from absl.testing import absltest +from absl.testing import parameterized import mujoco_warp as mjwarp from mujoco.mjx.third_party.mujoco_warp._src import test_util +from mujoco.mjx.third_party.mujoco_warp._src.io import MAX_WORLDS def _dims_match(test_obj, d1: Any, d2: Any, prefix: str = ""): @@ -173,7 +175,7 @@ def _leading_dims_scale_w_nworld(test_obj, d1: Any, d2: Any, nworld1: int, nworl test_obj.assertEqual(s1, nworld1, full_name + f" has leading dim {s1} with nworld={nworld1}. {msg}") -class IOTest(absltest.TestCase): +class IOTest(parameterized.TestCase): def test_make_put_data(self): """Tests that make_data and put_data are producing the same shapes for all arrays.""" mjm, _, _, d = test_util.fixture("pendula.xml") @@ -295,11 +297,11 @@ class IOTest(absltest.TestCase): m1 = mjwarp.put_model(mjm) self.assertTrue(hasattr(m1.geom_pos, "_is_batched")) - self.assertEqual(m1.geom_pos.shape[0], 1) + self.assertEqual(m1.geom_pos.shape[0], MAX_WORLDS) self.assertEqual(m1.geom_pos.strides[0], 0) self.assertLen(m1.geom_pos.strides, m1.geom_pos.ndim) self.assertTrue(hasattr(m1.opt.gravity, "_is_batched")) - self.assertEqual(m1.opt.gravity.shape[0], 1) + self.assertEqual(m1.opt.gravity.shape[0], MAX_WORLDS) self.assertEqual(m1.opt.gravity.strides[0], 0) self.assertLen(m1.opt.gravity.strides, m1.opt.gravity.ndim) self.assertFalse(hasattr(m1.body_parentid, "_is_batched")) @@ -348,6 +350,34 @@ class IOTest(absltest.TestCase): _dims_match(self, dm2, dp2) _dims_match(self, dm3, dp3) + @parameterized.parameters( + '', + '', + '', + '', + '', + '', + ) + def test_contact_sensor(self, contact_sensor): + mjm = mujoco.MjModel.from_xml_string(f""" + + + + + + + + + + + {contact_sensor} + + + """) + + with self.assertRaises(NotImplementedError): + mjwarp.put_model(mjm) + if __name__ == "__main__": wp.init() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/jax_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/jax_test.py index 529a2288..9965203c 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/jax_test.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/jax_test.py @@ -26,8 +26,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.test_util import fixture class JAXTest(parameterized.TestCase): - @parameterized.parameters("humanoid/humanoid.xml", "pendula.xml") - def test_jax(self, xml): + @parameterized.product(xml=("pendula.xml", "humanoid/humanoid.xml"), graph_conditional=(True, False)) + def test_jax(self, xml, graph_conditional): os.environ["XLA_FLAGS"] = "--xla_gpu_graph_min_graph_size=1" # Force JAX to allocate memory on demand and deallocate when not needed (slow) os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" @@ -38,7 +38,7 @@ class JAXTest(parameterized.TestCase): self.skipTest("JAX not installed") from jax import numpy as jp - from warp.jax_experimental.ffi import jax_callable + from mujoco.mjx.third_party.warp.jax_experimental import ffi if jax.default_backend() != "gpu": self.skipTest("JAX default backend is not GPU") @@ -51,14 +51,12 @@ class JAXTest(parameterized.TestCase): xml, nworld=NWORLDS, nconmax=NWORLDS * NCONTACTS, - njmax=NWORLDS * NCONTACTS * 4, + njmax=NCONTACTS * 4, iterations=1, ls_iterations=4, kick=True, ) - - # Disable CUDA graph conditional - m.opt.graph_conditional = False + m.opt.graph_conditional = graph_conditional def warp_step( qpos_in: wp.array(dtype=wp.float32, ndim=2), @@ -82,14 +80,16 @@ class JAXTest(parameterized.TestCase): return qpos, qvel - warp_step_fn = jax_callable( + warp_step_fn = ffi.jax_callable( warp_step, num_outputs=2, output_dims={"qpos_out": (NWORLDS, mjm.nq), "qvel_out": (NWORLDS, mjm.nv)}, - graph_compatible=True, + graph_mode=ffi.GraphMode.WARP, ) - jax_qpos = jp.tile(jp.array(m.qpos0.numpy()), (NWORLDS, 1)) + # temp qpos0 array to get the right numpy shape + qpos0_temp = wp.array(ptr=m.qpos0.ptr, shape=(1,) + m.qpos0.shape[1:], dtype=wp.float32) + jax_qpos = jp.tile(jp.array(qpos0_temp), (NWORLDS, 1)) jax_qvel = jp.zeros((NWORLDS, m.nv)) jax_unroll_fn = jax.jit(unroll).lower(jax_qpos, jax_qvel).compile() 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 ecf5f322..48175829 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -21,17 +21,21 @@ from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import ray from mujoco.mjx.third_party.mujoco_warp._src import smooth from mujoco.mjx.third_party.mujoco_warp._src import support +from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import DataType 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 JointType from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType +from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6 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 @@ -1546,6 +1550,7 @@ def _frameangacc( @wp.kernel def _sensor_acc( # Model: + opt_cone: int, body_rootid: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), @@ -1555,10 +1560,15 @@ def _sensor_acc( sensor_datatype: wp.array(dtype=int), sensor_objtype: wp.array(dtype=int), sensor_objid: wp.array(dtype=int), + sensor_intprm: wp.array2d(dtype=int), + sensor_dim: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), sensor_acc_adr: wp.array(dtype=int), + sensor_contact_adr: wp.array(dtype=int), # Data in: + njmax_in: int, + ncon_in: wp.array(dtype=int), xpos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1569,8 +1579,18 @@ def _sensor_acc( cvel_in: wp.array2d(dtype=wp.spatial_vector), actuator_force_in: wp.array2d(dtype=float), qfrc_actuator_in: wp.array2d(dtype=float), + contact_dist_in: wp.array(dtype=float), + contact_pos_in: wp.array(dtype=wp.vec3), + contact_frame_in: wp.array(dtype=wp.mat33), + contact_friction_in: wp.array(dtype=vec5), + contact_dim_in: wp.array(dtype=int), + contact_efc_address_in: wp.array2d(dtype=int), + efc_force_in: wp.array2d(dtype=float), cacc_in: wp.array2d(dtype=wp.spatial_vector), cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + sensor_contact_nmatch_in: wp.array2d(dtype=int), + sensor_contact_matchid_in: wp.array3d(dtype=int), + sensor_contact_direction_in: wp.array3d(dtype=float), # Data out: sensordata_out: wp.array2d(dtype=float), ): @@ -1580,7 +1600,123 @@ def _sensor_acc( objid = sensor_objid[sensorid] out = sensordata_out[worldid] - if sensortype == int(SensorType.ACCELEROMETER.value): + if sensortype == int(SensorType.CONTACT.value): + dataspec = sensor_intprm[sensorid, 0] + dim = sensor_dim[sensorid] + objtype = sensor_objtype[sensorid] + + # found, force, torque, dist, pos, normal, tangent + # TODO(thowell): precompute slot size + found = False + force = False + torque = False + dist = False + pos = False + normal = False + tangent = False + + size = int(0) + for i in range(7): + if dataspec & (1 << i): + if i == 0: + found = True + size += 1 + elif i == 1: + force = True + size += 3 + elif i == 2: + torque = True + size += 3 + elif i == 3: + dist = True + size += 1 + elif i == 4: + pos = True + size += 3 + elif i == 5: + normal = True + size += 3 + elif i == 6: + tangent = True + size += 3 + + num = dim // size # number of slots + + adr = sensor_adr[sensorid] + + # TODO(team): precompute sensorid to contactsensorid mapping + contactsensorid = int(0) + for i in range(sensor_contact_adr.size): + if sensorid == sensor_contact_adr[i]: + contactsensorid = i + break + + nmatch = sensor_contact_nmatch_in[worldid, contactsensorid] + + for i in range(wp.min(nmatch, num)): + # sorted contact id + cid = sensor_contact_matchid_in[worldid, contactsensorid, i] + + # contact direction + dir = sensor_contact_direction_in[worldid, contactsensorid, i] + + adr_slot = adr + i * size + + if found: + out[adr_slot] = float(nmatch) + adr_slot += 1 + if force or torque: + contact_forcetorque = 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, + cid, + False, + ) + if force: + out[adr_slot + 0] = contact_forcetorque[0] + out[adr_slot + 1] = contact_forcetorque[1] + out[adr_slot + 2] = dir * contact_forcetorque[2] + adr_slot += 3 + if torque: + out[adr_slot + 0] = contact_forcetorque[3] + out[adr_slot + 1] = contact_forcetorque[4] + out[adr_slot + 2] = dir * contact_forcetorque[5] + adr_slot += 3 + if dist: + out[adr_slot] = contact_dist_in[cid] + adr_slot += 1 + if pos: + contact_pos = contact_pos_in[cid] + out[adr_slot + 0] = contact_pos[0] + out[adr_slot + 1] = contact_pos[1] + out[adr_slot + 2] = contact_pos[2] + adr_slot += 3 + if normal: + contact_normal = contact_frame_in[cid][0] + out[adr_slot + 0] = dir * contact_normal[0] + out[adr_slot + 1] = dir * contact_normal[1] + out[adr_slot + 2] = dir * contact_normal[2] + adr_slot += 3 + if tangent: + contact_tangent = contact_frame_in[cid][1] + out[adr_slot + 0] = dir * contact_tangent[0] + out[adr_slot + 1] = dir * contact_tangent[1] + out[adr_slot + 2] = dir * contact_tangent[2] + adr_slot += 3 + + # zero remaining slots + for i in range(nmatch, num): + for j in range(size): + out[adr + i * size + j] = 0.0 + + elif sensortype == int(SensorType.ACCELEROMETER.value): vec3 = _accelerometer( body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid ) @@ -1726,6 +1862,210 @@ def _sensor_touch( wp.atomic_add(sensordata_out[worldid], adr, normalforce) +@wp.kernel +def _sensor_tactile_zero( + # Model: + sensor_type: wp.array(dtype=int), + sensor_dim: wp.array(dtype=int), + sensor_adr: wp.array(dtype=int), + # Data out: + sensordata_out: wp.array2d(dtype=float), +): + worldid, sensorid = wp.tid() + + if sensor_type[sensorid] != int(SensorType.TACTILE.value): + return + + for i in range(sensor_dim[sensorid]): + sensordata_out[worldid, sensor_adr[sensorid] + i] = 0.0 + + +@wp.func +def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3: + return wp.spatial_bottom(vec) - wp.cross(dif, wp.spatial_top(vec)) + + +@wp.kernel +def _sensor_tactile( + # Model: + body_rootid: wp.array(dtype=int), + body_weldid: wp.array(dtype=int), + geom_bodyid: wp.array(dtype=int), + mesh_vertadr: wp.array(dtype=int), + mesh_vert: wp.array(dtype=wp.vec3), + mesh_normaladr: wp.array(dtype=int), + mesh_normal: wp.array(dtype=wp.vec3), + mesh_quat: wp.array(dtype=wp.quat), + sensor_objid: wp.array(dtype=int), + sensor_refid: wp.array(dtype=int), + sensor_dim: wp.array(dtype=int), + sensor_adr: wp.array(dtype=int), + plugin: wp.array(dtype=int), + plugin_attr: wp.array(dtype=wp.vec3f), + geom_plugin_index: wp.array(dtype=int), + taxel_vertadr: wp.array(dtype=int), + taxel_sensorid: wp.array(dtype=int), + # Data in: + ncon_in: wp.array(dtype=int), + geom_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xmat_in: wp.array2d(dtype=wp.mat33), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cvel_in: wp.array2d(dtype=wp.spatial_vector), + contact_geom_in: wp.array(dtype=wp.vec2i), + contact_worldid_in: wp.array(dtype=int), + # Data out: + sensordata_out: wp.array2d(dtype=float), +): + conid, taxelid = wp.tid() + + if conid >= ncon_in[0]: + return + + worldid = contact_worldid_in[conid] + + # get sensor_id + sensor_id = taxel_sensorid[taxelid] + + # get parent weld id + mesh_id = sensor_objid[sensor_id] + geom_id = sensor_refid[sensor_id] + parent_body = geom_bodyid[geom_id] + parent_weld = body_weldid[parent_body] + + # contact geom + body1 = body_weldid[geom_bodyid[contact_geom_in[conid][0]]] + body2 = body_weldid[geom_bodyid[contact_geom_in[conid][1]]] + if body1 == parent_weld: + geom = contact_geom_in[conid][1] + elif body2 == parent_weld: + geom = contact_geom_in[conid][0] + else: + return + body = geom_bodyid[geom] + + # vertex local position + vertid = taxel_vertadr[taxelid] - mesh_vertadr[mesh_id] + pos = mesh_vert[vertid + mesh_vertadr[mesh_id]] + + # position in global frame + xpos = geom_xmat_in[worldid, geom_id] @ pos + xpos += geom_xpos_in[worldid, geom_id] + + # position in other geom frame + tmp = xpos - geom_xpos_in[worldid, geom] + lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp + + # compute distance + plugin_id = geom_plugin_index[geom] + depth = wp.min(sdf(int(GeomType.SDF.value), lpos, plugin_attr[plugin_id], plugin[plugin_id]), 0.0) + if depth >= 0.0: + return + + # get velocity in global + vel_sensor = _transform_spatial(cvel_in[worldid, parent_weld], xpos - subtree_com_in[worldid, body_rootid[parent_weld]]) + vel_other = _transform_spatial( + cvel_in[worldid, body], geom_xpos_in[worldid, geom] - subtree_com_in[worldid, body_rootid[body]] + ) + vel_rel = vel_sensor - vel_other + + # get contact force/torque, rotate into node frame + offset = mesh_normaladr[mesh_id] + 3 * vertid + normal = math.rot_vec_quat(mesh_normal[offset], mesh_quat[mesh_id]) + tang1 = math.rot_vec_quat(mesh_normal[offset + 1], mesh_quat[mesh_id]) + tang2 = math.rot_vec_quat(mesh_normal[offset + 2], mesh_quat[mesh_id]) + kMaxDepth = 0.05 + pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL) + force = wp.mul(normal, pressure) + + # one row of mat^T * force + forceT = wp.vec3() + forceT[0] = wp.dot(force, normal) + forceT[1] = wp.abs(wp.dot(vel_rel, tang1)) + forceT[2] = wp.abs(wp.dot(vel_rel, tang2)) + + # add to sensor output + dim = sensor_dim[sensor_id] / 3 + wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0]) + wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1]) + wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2]) + + +@wp.kernel +def _contact_match( + # Model: + sensor_objid: wp.array(dtype=int), + sensor_refid: wp.array(dtype=int), + sensor_contact_adr: wp.array(dtype=int), + # Data in: + ncon_in: wp.array(dtype=int), + contact_dist_in: wp.array(dtype=float), + contact_geom_in: wp.array(dtype=wp.vec2i), + contact_worldid_in: wp.array(dtype=int), + # Data out: + sensor_contact_nmatch_out: wp.array2d(dtype=int), + sensor_contact_matchid_out: wp.array3d(dtype=int), + sensor_contact_criteria_out: wp.array3d(dtype=float), + sensor_contact_direction_out: wp.array3d(dtype=float), +): + contactsensorid, contactid = wp.tid() + sensorid = sensor_contact_adr[contactsensorid] + + if contactid >= ncon_in[0]: + return + + # sensor information + objid = sensor_objid[sensorid] + refid = sensor_refid[sensorid] + + # contact information + geom = contact_geom_in[contactid] + + # geom-geom match + geom0geom1 = objid == geom[0] and refid == geom[1] + geom1geom0 = objid == geom[1] and refid == geom[0] + if geom0geom1 or geom1geom0: + worldid = contact_worldid_in[contactid] + + 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] + + # contact direction + if geom1geom0: + sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = -1.0 + else: + sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = 1.0 + + return + + # TODO(thowell): alternative matching + + +@wp.kernel +def _contact_sort( + # Data in: + sensor_contact_nmatch_in: wp.array2d(dtype=int), + sensor_contact_matchid_in: wp.array3d(dtype=int), + sensor_contact_criteria_in: wp.array3d(dtype=float), + # Data out: + sensor_contact_matchid_out: wp.array3d(dtype=int), +): + worldid, contactsensorid = wp.tid() + + nmatch = sensor_contact_nmatch_in[worldid, contactsensorid] + + # skip sort + if nmatch <= 1: + return + + criteria_tile = wp.tile_load(sensor_contact_criteria_in[worldid, contactsensorid], shape=MJ_MAXCONPAIR) + matchid_tile = wp.tile_load(sensor_contact_matchid_in[worldid, contactsensorid], shape=MJ_MAXCONPAIR) + wp.tile_sort(criteria_tile, matchid_tile) + wp.tile_store(sensor_contact_matchid_out[worldid, contactsensorid], matchid_tile) + + @event_scope def sensor_acc(m: Model, d: Data): """Compute acceleration-dependent sensor values.""" @@ -1772,6 +2112,94 @@ def sensor_acc(m: Model, d: Data): ], ) + wp.launch( + _sensor_tactile_zero, + dim=(d.nworld, m.nsensordata), + inputs=[ + m.sensor_type, + m.sensor_dim, + m.sensor_adr, + ], + outputs=[ + d.sensordata, + ], + ) + + wp.launch( + _sensor_tactile, + dim=(d.nconmax, m.nsensortaxel), + inputs=[ + m.body_rootid, + m.body_weldid, + m.geom_bodyid, + m.mesh_vertadr, + m.mesh_vert, + m.mesh_normaladr, + m.mesh_normal, + m.mesh_quat, + m.sensor_objid, + m.sensor_refid, + m.sensor_dim, + m.sensor_adr, + m.plugin, + m.plugin_attr, + m.geom_plugin_index, + m.taxel_vertadr, + m.taxel_sensorid, + d.ncon, + d.geom_xpos, + d.geom_xmat, + d.subtree_com, + d.cvel, + d.contact.geom, + d.contact.worldid, + ], + outputs=[ + d.sensordata, + ], + ) + + if m.sensor_contact_adr.size: + # match criteria + d.sensor_contact_nmatch.zero_() + d.sensor_contact_matchid.zero_() + d.sensor_contact_criteria.zero_() + + wp.launch( + _contact_match, + dim=(m.sensor_contact_adr.size, d.nconmax), + inputs=[ + m.sensor_objid, + m.sensor_refid, + m.sensor_contact_adr, + d.ncon, + d.contact.dist, + d.contact.geom, + d.contact.worldid, + ], + outputs=[ + d.sensor_contact_nmatch, + d.sensor_contact_matchid, + d.sensor_contact_criteria, + d.sensor_contact_direction, + ], + ) + + # sorting + wp.launch_tiled( + _contact_sort, + dim=(d.nworld, m.sensor_contact_adr.size), + inputs=[ + d.sensor_contact_nmatch, + d.sensor_contact_matchid, + d.sensor_contact_criteria, + ], + outputs=[ + d.sensor_contact_matchid, + ], + block_dim=m.block_dim.contact_sort, + ) + if m.sensor_rne_postconstraint: smooth.rne_postconstraint(m, d) @@ -1779,6 +2207,7 @@ def sensor_acc(m: Model, d: Data): _sensor_acc, dim=(d.nworld, m.sensor_acc_adr.size), inputs=[ + m.opt.cone, m.body_rootid, m.jnt_dofadr, m.geom_bodyid, @@ -1788,9 +2217,14 @@ def sensor_acc(m: Model, d: Data): m.sensor_datatype, m.sensor_objtype, m.sensor_objid, + m.sensor_intprm, + m.sensor_dim, m.sensor_adr, m.sensor_cutoff, m.sensor_acc_adr, + m.sensor_contact_adr, + d.njmax, + d.ncon, d.xpos, d.xipos, d.geom_xpos, @@ -1801,8 +2235,18 @@ def sensor_acc(m: Model, d: Data): d.cvel, d.actuator_force, d.qfrc_actuator, + d.contact.dist, + d.contact.pos, + d.contact.frame, + d.contact.friction, + d.contact.dim, + d.contact.efc_address, + d.efc.force, d.cacc, d.cfrc_int, + d.sensor_contact_nmatch, + d.sensor_contact_matchid, + d.sensor_contact_direction, ], outputs=[d.sensordata], ) 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 458dea1f..ff36df1c 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 @@ -15,6 +15,8 @@ """Tests for sensor functions.""" +import itertools + import mujoco import numpy as np import warp as wp @@ -405,6 +407,66 @@ class SensorTest(parameterized.TestCase): _assert_eq(d.energy.numpy()[0][1], mjd.energy[1], "kinetic energy") + @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"', + ) + def test_contact_sensor(self, geom): + """Test contact sensor.""" + # create contact sensors + 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))]) + + 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'' + + _MJCF = f""" + + + + """ + + _, mjd, m, d = test_util.fixture(xml=_MJCF, keyframe=0) + + d.sensordata.zero_() + mjwarp.forward(m, d) + + sensordata = d.sensordata.numpy()[0] + _assert_eq(sensordata, mjd.sensordata, "sensordata") + self.assertTrue(sensordata.any()) # check that sensordata is not empty + if __name__ == "__main__": wp.init() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py index b22ee540..73276fc3 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py @@ -1252,6 +1252,7 @@ def _cfrc_ext_contact( body_rootid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), # Data in: + njmax_in: int, ncon_in: wp.array(dtype=int), subtree_com_in: wp.array2d(dtype=wp.vec3), contact_pos_in: wp.array(dtype=wp.vec3), @@ -1282,6 +1283,7 @@ def _cfrc_ext_contact( # contact force in world frame force = support.contact_force_fn( opt_cone, + njmax_in, ncon_in, contact_frame_in, contact_friction_in, @@ -1351,6 +1353,7 @@ def rne_postconstraint(m: Model, d: Data): m.opt.cone, m.body_rootid, m.geom_bodyid, + d.njmax, d.ncon, d.subtree_com, d.contact.pos, 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 31de24df..ed74254f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -14,6 +14,7 @@ # ============================================================================== from math import ceil +from math import sqrt import warp as wp @@ -1651,8 +1652,6 @@ def update_constraint_zero_qfrc_constraint( @wp.kernel def update_constraint_init_qfrc_constraint( - # Model: - nv: int, # Data in: nefc_in: wp.array(dtype=int), efc_J_in: wp.array3d(dtype=float), @@ -1660,48 +1659,59 @@ def update_constraint_init_qfrc_constraint( efc_done_in: wp.array(dtype=bool), # Data out: qfrc_constraint_out: wp.array2d(dtype=float), -): - worldid, efcid = wp.tid() - - if efcid >= nefc_in[worldid]: - return - - if efc_done_in[worldid]: - return - - force = efc_force_in[worldid, efcid] - for i in range(nv): - wp.atomic_add( - qfrc_constraint_out[worldid], - i, - efc_J_in[worldid, efcid, i] * force, - ) - - -@wp.kernel -def update_constraint_gauss_cost( - # Data in: - qacc_in: wp.array2d(dtype=float), - qfrc_smooth_in: wp.array2d(dtype=float), - qacc_smooth_in: wp.array2d(dtype=float), - efc_Ma_in: wp.array2d(dtype=float), - efc_done_in: wp.array(dtype=bool), - # Data out: - efc_gauss_out: wp.array(dtype=float), - efc_cost_out: wp.array(dtype=float), ): worldid, dofid = wp.tid() if efc_done_in[worldid]: return - gauss_cost = ( - 0.5 - * (efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid]) - * (qacc_in[worldid, dofid] - qacc_smooth_in[worldid, dofid]) - ) - wp.atomic_add(efc_gauss_out, worldid, gauss_cost) - wp.atomic_add(efc_cost_out, worldid, gauss_cost) + sum_qfrc = float(0.0) + for efcid in range(nefc_in[worldid]): + efc_J = efc_J_in[worldid, efcid, dofid] + force = efc_force_in[worldid, efcid] + sum_qfrc += efc_J * force + + qfrc_constraint_out[worldid, dofid] += sum_qfrc + + +@cache_kernel +def update_constraint_gauss_cost(nv: int, dofs_per_thread: int): + @nested_kernel + def kernel( + # Data in: + qacc_in: wp.array2d(dtype=float), + qfrc_smooth_in: wp.array2d(dtype=float), + qacc_smooth_in: wp.array2d(dtype=float), + efc_Ma_in: wp.array2d(dtype=float), + efc_done_in: wp.array(dtype=bool), + # Data out: + efc_gauss_out: wp.array(dtype=float), + efc_cost_out: wp.array(dtype=float), + ): + worldid, dofstart = wp.tid() + + if efc_done_in[worldid]: + return + + gauss_cost = float(0.0) + + if wp.static(dofs_per_thread >= nv): + for i in range(wp.static(min(dofs_per_thread, nv))): + gauss_cost += (efc_Ma_in[worldid, i] - qfrc_smooth_in[worldid, i]) * (qacc_in[worldid, i] - qacc_smooth_in[worldid, i]) + efc_gauss_out[worldid] += 0.5 * gauss_cost + efc_cost_out[worldid] += 0.5 * gauss_cost + + else: + for i in range(wp.static(dofs_per_thread)): + ii = dofstart * wp.static(dofs_per_thread) + i + if ii < nv: + gauss_cost += (efc_Ma_in[worldid, ii] - qfrc_smooth_in[worldid, ii]) * ( + qacc_in[worldid, ii] - qacc_smooth_in[worldid, ii] + ) + wp.atomic_add(efc_gauss_out, worldid, gauss_cost) + wp.atomic_add(efc_cost_out, worldid, gauss_cost) + + return kernel def _update_constraint(m: types.Model, d: types.Data): @@ -1813,15 +1823,24 @@ def _update_constraint(m: types.Model, d: types.Data): wp.launch( update_constraint_init_qfrc_constraint, - dim=(d.nworld, d.njmax), - inputs=[m.nv, d.nefc, d.efc.J, d.efc.force, d.efc.done], + dim=(d.nworld, m.nv), + inputs=[d.nefc, d.efc.J, d.efc.force, d.efc.done], outputs=[d.qfrc_constraint], ) + # if we are only using 1 thread, it makes sense to do more dofs and skip the atomics. + # For more than 1 thread, dofs_per_thread is lower for better load balancing. + if m.nv > 50: + dofs_per_thread = 20 + else: + dofs_per_thread = 50 + + threads_per_efc = ceil(m.nv / dofs_per_thread) + # gauss = 0.5 * (Ma - qfrc_smooth).T @ (qacc - qacc_smooth) wp.launch( - update_constraint_gauss_cost, - dim=(d.nworld, m.nv), + update_constraint_gauss_cost(m.nv, dofs_per_thread), + dim=(d.nworld, threads_per_efc), inputs=[d.qacc, d.qfrc_smooth, d.qacc_smooth, d.efc.Ma, d.efc.done], outputs=[d.efc.gauss, d.efc.cost], ) @@ -1927,16 +1946,12 @@ def update_gradient_copy_lower_triangle( @wp.kernel def update_gradient_JTDAJ( - # Model: - dof_tri_row: wp.array(dtype=int), - dof_tri_col: wp.array(dtype=int), # Data in: nefc_in: wp.array(dtype=int), efc_J_in: wp.array3d(dtype=float), efc_D_in: wp.array2d(dtype=float), efc_active_in: wp.array2d(dtype=bool), efc_done_in: wp.array(dtype=bool), - # In: # Data out: efc_h_out: wp.array3d(dtype=float), ): @@ -1947,20 +1962,26 @@ def update_gradient_JTDAJ( nefc = nefc_in[worldid] - dofi = dof_tri_row[elementid] - dofj = dof_tri_col[elementid] - - for efcid in range(nefc): - efc_D = efc_D_in[worldid, efcid] - active = efc_active_in[worldid, efcid] - - if efc_D == 0.0 or not active: - continue + dofi = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2 + dofj = elementid - (dofi * (dofi + 1)) // 2 + sum_h = float(0.0) + efc_D = efc_D_in[worldid, 0] + active = efc_active_in[worldid, 0] + efc_Ji = efc_J_in[worldid, 0, dofi] + efc_Jj = efc_J_in[worldid, 0, dofj] + for efcid in range(nefc - 1): # TODO(team): sparse efc_J - value = efc_J_in[worldid, efcid, dofi] * efc_J_in[worldid, efcid, dofj] * efc_D - if value != 0.0: - wp.atomic_add(efc_h_out[worldid, dofi], dofj, value) + sum_h += efc_Ji * efc_Jj * efc_D * float(active) + + jj = efcid + 1 + efc_D = efc_D_in[worldid, jj] + active = efc_active_in[worldid, jj] + efc_Ji = efc_J_in[worldid, jj, dofi] + efc_Jj = efc_J_in[worldid, jj, dofj] + + sum_h += efc_Ji * efc_Jj * efc_D * float(active) + efc_h_out[worldid, dofi, dofj] += sum_h @wp.kernel @@ -2184,12 +2205,12 @@ def _update_gradient(m: types.Model, d: types.Data): 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, m.dof_tri_row.size), + dim=(d.nworld, lower_triangle_dim), inputs=[ - m.dof_tri_row, - m.dof_tri_col, d.nefc, d.efc.J, d.efc.D, @@ -2295,24 +2316,9 @@ def solve_prev_grad_Mgrad( @wp.kernel -def solve_zero_beta_num_den( - # Data in: - efc_done_in: wp.array(dtype=bool), - # Data out: - efc_beta_num_out: wp.array(dtype=float), - efc_beta_den_out: wp.array(dtype=float), -): - worldid = wp.tid() - - if efc_done_in[worldid]: - return - - efc_beta_num_out[worldid] = 0.0 - efc_beta_den_out[worldid] = 0.0 - - -@wp.kernel -def solve_beta_num_den( +def solve_beta( + # Model: + nv: int, # Data in: efc_grad_in: wp.array2d(dtype=float), efc_Mgrad_in: wp.array2d(dtype=float), @@ -2320,30 +2326,6 @@ def solve_beta_num_den( efc_prev_Mgrad_in: wp.array2d(dtype=float), efc_done_in: wp.array(dtype=bool), # Data out: - efc_beta_num_out: wp.array(dtype=float), - efc_beta_den_out: wp.array(dtype=float), -): - worldid, dofid = wp.tid() - - if efc_done_in[worldid]: - return - - prev_Mgrad = efc_prev_Mgrad_in[worldid][dofid] - wp.atomic_add( - efc_beta_num_out, - worldid, - efc_grad_in[worldid, dofid] * (efc_Mgrad_in[worldid, dofid] - prev_Mgrad), - ) - wp.atomic_add(efc_beta_den_out, worldid, efc_prev_grad_in[worldid, dofid] * prev_Mgrad) - - -@wp.kernel -def solve_beta( - # Data in: - efc_beta_num_in: wp.array(dtype=float), - efc_beta_den_in: wp.array(dtype=float), - efc_done_in: wp.array(dtype=bool), - # Data out: efc_beta_out: wp.array(dtype=float), ): worldid = wp.tid() @@ -2351,7 +2333,14 @@ def solve_beta( if efc_done_in[worldid]: return - efc_beta_out[worldid] = wp.max(0.0, efc_beta_num_in[worldid] / wp.max(types.MJ_MINVAL, efc_beta_den_in[worldid])) + beta_num = float(0.0) + beta_den = float(0.0) + for dofid in range(nv): + prev_Mgrad = efc_prev_Mgrad_in[worldid][dofid] + beta_num += efc_grad_in[worldid, dofid] * (efc_Mgrad_in[worldid, dofid] - prev_Mgrad) + beta_den += efc_prev_grad_in[worldid, dofid] * prev_Mgrad + + efc_beta_out[worldid] = wp.max(0.0, beta_num / wp.max(types.MJ_MINVAL, beta_den)) @wp.kernel @@ -2451,24 +2440,10 @@ def _solver_iteration( # polak-ribiere if m.opt.solver == types.SolverType.CG: - wp.launch( - solve_zero_beta_num_den, - dim=(d.nworld), - inputs=[d.efc.done], - outputs=[d.efc.beta_num, d.efc.beta_den], - ) - - wp.launch( - solve_beta_num_den, - dim=(d.nworld, m.nv), - inputs=[d.efc.grad, d.efc.Mgrad, d.efc.prev_grad, d.efc.prev_Mgrad, d.efc.done], - outputs=[d.efc.beta_num, d.efc.beta_den], - ) - wp.launch( solve_beta, dim=(d.nworld,), - inputs=[d.efc.beta_num, d.efc.beta_den, d.efc.done], + inputs=[m.nv, d.efc.grad, d.efc.Mgrad, d.efc.prev_grad, d.efc.prev_Mgrad, d.efc.done], outputs=[d.efc.beta], ) @@ -2539,8 +2514,10 @@ def solve(m: types.Model, d: types.Data): def _solve(m: types.Model, d: types.Data): """Finds forces that satisfy constraints.""" - # warmstart - wp.copy(d.qacc, d.qacc_warmstart) + if not (m.opt.disableflags & types.DisableBit.WARMSTART): + wp.copy(d.qacc, d.qacc_warmstart) + else: + wp.copy(d.qacc, d.qacc_smooth) # create context create_context(m, d, grad=True) 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 11d043f6..c4629d20 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 @@ -68,6 +68,63 @@ class SolverTest(parameterized.TestCase): _assert_eq(mjwarp_cost, mj_cost, name="cost") + @parameterized.parameters(ConeType.PYRAMIDAL, ConeType.ELLIPTIC) + 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, + ls_parallel=False, + iterations=50, + ls_iterations=50, + ) + + # One step to obtain more non-zeros results + mjwarp.step(m, d) + + # Preparing for linesearch + m.opt.iterations = 0 + mjwarp.fwd_velocity(m, d) + mjwarp.fwd_acceleration(m, d, factorize=True) + solver.solve(m, d) + + # Storing some initial values + d_efc_Ma = d.efc.Ma.numpy().copy() + d_efc_Jaref = d.efc.Jaref.numpy().copy() + d_qacc = d.qacc.numpy().copy() + + # Launching iterative linesearch + m.opt.ls_parallel = False + solver._linesearch(m, d) + alpha_iterative = d.efc.alpha.numpy().copy() + + # Launching parallel linesearch with 10 testing points + m.nlsp = 10 + d.efc.Ma = wp.array2d(d_efc_Ma) + d.efc.Jaref = wp.array(d_efc_Jaref) + d.qacc = wp.array2d(d_qacc) + m.opt.ls_parallel = True + solver._linesearch(m, d) + alpha_parallel_10 = d.efc.alpha.numpy().copy() + + # Launching parallel linesearch with 50 testing points + m.nlsp = 50 + d.efc.Ma = wp.array2d(d_efc_Ma) + d.efc.Jaref = wp.array(d_efc_Jaref) + d.qacc = wp.array2d(d_qacc) + solver._linesearch(m, d) + alpha_parallel_50 = d.efc.alpha.numpy().copy() + + # Checking that iterative and parallel linesearch lead to similar results + # and that increasing ls_iterations leads to better results + _assert_eq(alpha_iterative, alpha_parallel_50, name="linesearch alpha") + self.assertLessEqual(abs(alpha_iterative - alpha_parallel_50), abs(alpha_iterative - alpha_parallel_10)) + @parameterized.parameters( (ConeType.PYRAMIDAL, SolverType.CG, 5, 5, False, False), (ConeType.ELLIPTIC, SolverType.CG, 5, 5, False, False), 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 b59ffadd..5d276fbe 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -298,7 +298,9 @@ def any_different(v0: wp.vec3, v1: wp.vec3) -> wp.bool: @wp.func -def _decode_pyramid(pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int) -> wp.spatial_vector: +def _decode_pyramid( + njmax_in: int, pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int +) -> wp.spatial_vector: """Converts pyramid representation to contact force.""" force = wp.spatial_vector() @@ -308,8 +310,15 @@ def _decode_pyramid(pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, force[0] = float(0.0) for i in range(condim - 1): - dir1 = pyramid[2 * i + efc_address] - dir2 = pyramid[2 * i + efc_address + 1] + adr = 2 * i + efc_address + if adr < njmax_in: + dir1 = pyramid[adr] + else: + dir1 = 0.0 + if adr + 1 < njmax_in: + dir2 = pyramid[adr + 1] + else: + dir2 = 0.0 force[0] += dir1 + dir2 force[i + 1] = (dir1 - dir2) * mu[i] @@ -321,6 +330,7 @@ def contact_force_fn( # Model: opt_cone: int, # Data in: + njmax_in: int, ncon_in: wp.array(dtype=int), contact_frame_in: wp.array(dtype=wp.mat33), contact_friction_in: wp.array(dtype=vec5), @@ -340,6 +350,7 @@ def contact_force_fn( if contact_id >= 0 and contact_id <= ncon_in[0] and efc_address >= 0: if opt_cone == int(ConeType.PYRAMIDAL.value): force = _decode_pyramid( + njmax_in, efc_force_in[worldid], efc_address, contact_friction_in[contact_id], @@ -347,7 +358,8 @@ def contact_force_fn( ) else: for i in range(condim): - force[i] = efc_force_in[worldid, contact_efc_address_in[contact_id, i]] + if contact_efc_address_in[contact_id, i] < njmax_in: + force[i] = efc_force_in[worldid, contact_efc_address_in[contact_id, i]] if to_world_frame: # Transform both top and bottom parts of spatial vector by the full contact frame matrix @@ -363,6 +375,7 @@ def contact_force_kernel( # Model: opt_cone: int, # Data in: + njmax_in: int, ncon_in: wp.array(dtype=int), contact_frame_in: wp.array(dtype=wp.mat33), contact_friction_in: wp.array(dtype=vec5), @@ -387,6 +400,7 @@ def contact_force_kernel( out[tid] = contact_force_fn( opt_cone, + njmax_in, ncon_in, contact_frame_in, contact_friction_in, @@ -421,6 +435,7 @@ def contact_force( dim=(contact_ids.size,), inputs=[ m.opt.cone, + d.njmax, d.ncon, d.contact.frame, d.contact.friction, 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 4d9ae7b1..b78de634 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 @@ -15,6 +15,8 @@ """Utilities for testing.""" +import importlib +import os import time from typing import Callable, Optional, Tuple @@ -23,6 +25,7 @@ import numpy as np import warp as wp from etils import epath +from mujoco.mjx.third_party.mujoco_warp._src import forward from mujoco.mjx.third_party.mujoco_warp._src import io from mujoco.mjx.third_party.mujoco_warp._src import warp_util from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType @@ -242,26 +245,107 @@ def benchmark( m.actuator_ctrllimited, m.actuator_ctrlrange, i, 0.01 ], outputs=[d.ctrl]) # fmt: skip + wp.synchronize() run_beg = time.perf_counter() wp.capture_launch(graph) wp.synchronize() + run_end = time.perf_counter() - run_end = time.perf_counter() time_vec[i] = run_end - run_beg if trace: trace = _sum(trace, tracer.trace()) else: trace = tracer.trace() - if measure_alloc or measure_solver_niter: - wp.synchronize() if measure_alloc: ncon.append(d.ncon.numpy()[0]) nefc.append(np.sum(d.nefc.numpy())) if measure_solver_niter: solver_niter.append(d.solver_niter.numpy()) - wp.synchronize() run_duration = np.sum(time_vec) return jit_duration, run_duration, trace, ncon, nefc, solver_niter + + +class BenchmarkSuite: + """Base suite for all model benchmarks.""" + + path = "" + batch_size = -1 + nconmax = -1 + njmax = -1 + param_names = ("function",) + params = ( + "jit_duration", + "solver_niter_mean", + "solver_niter_p95", + "device_memory_allocated", + "step", + "step.forward", + "step.forward.fwd_position", + "step.forward.fwd_position.kinematics", + "step.forward.fwd_position.com_pos", + "step.forward.fwd_position.camlight", + "step.forward.fwd_position.crb", + "step.forward.fwd_position.tendon_armature", + "step.forward.fwd_position.collision", + "step.forward.fwd_position.make_constraint", + "step.forward.fwd_position.transmission", + "step.forward.sensor_pos", + "step.forward.fwd_velocity", + "step.forward.fwd_velocity.com_vel", + "step.forward.fwd_velocity.passive", + "step.forward.fwd_velocity.rne", + "step.forward.fwd_velocity.tendon_bias", + "step.forward.sensor_vel", + "step.forward.fwd_actuation", + "step.forward.fwd_acceleration", + "step.forward.fwd_acceleration.xfrc_accumulate", + "step.forward.sensor_acc", + "step.forward.solve", + ) + number = 1 + rounds = 1 + sample_time = 0 + repeat = 1 + + def setup_cache(self): + module = importlib.import_module(self.__module__) + path = os.path.join(os.path.realpath(os.path.dirname(module.__file__)), self.path) + mjm = mujoco.MjModel.from_xml_path(path) + mjd = mujoco.MjData(mjm) + if mjm.nkey > 0: + mujoco.mj_resetDataKeyframe(mjm, mjd, 0) + + # TODO(team): mj_forward call shouldn't be necessary, but it is + mujoco.mj_forward(mjm, mjd) + + wp.init() + + 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) + + 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, + } + + def tree_flatten(d, parent_k=""): + ret = {} + steps = self.batch_size * 1000 + for k, v in d.items(): + k = parent_k + "." + k if parent_k else k + ret = ret | {k: 1e6 * v[0][0] / steps} | tree_flatten(v[1], k) + return ret + + metrics = metrics | tree_flatten(trace) + + return metrics + + def track_metric(self, metrics, fn): + return metrics[fn] 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 f1e08dca..8c5197ee 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -48,6 +48,7 @@ class BlockDim: # ray ray: int = 64 # sensor + contact_sort: int = 64 energy_vel_kinetic: int = 256 # smooth cholesky_factorize: int = 256 @@ -59,6 +60,20 @@ class BlockDim: mul_m_dense: int = 256 +class BroadphaseType(enum.IntEnum): + """Type of broadphase algorithm. + + Attributes: + NXN: Broad phase checking all pairs + SAP_TILE: Sweep and prune broad phase using tile sort + SAP_SEGMENTED: Sweep and prune broad phase using segment sort + """ + + NXN = 0 + SAP_TILE = 1 + SAP_SEGMENTED = 2 + + class BroadphaseFilter(enum.IntFlag): """Bitmask specifying which collision functions to run during broadphase. @@ -133,12 +148,13 @@ class DisableBit(enum.IntFlag): PASSIVE = mujoco.mjtDisableBit.mjDSBL_PASSIVE GRAVITY = mujoco.mjtDisableBit.mjDSBL_GRAVITY CLAMPCTRL = mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL + WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART ACTUATION = mujoco.mjtDisableBit.mjDSBL_ACTUATION REFSAFE = mujoco.mjtDisableBit.mjDSBL_REFSAFE EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP FILTERPARENT = mujoco.mjtDisableBit.mjDSBL_FILTERPARENT SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR - # unsupported: MIDPHASE, WARMSTART + # unsupported: MIDPHASE class EnableBit(enum.IntFlag): @@ -371,6 +387,7 @@ class SensorType(enum.IntEnum): SUBTREELINVEL: subtree linear velocity SUBTREEANGMOM: subtree angular momentum TOUCH: scalar contact normal forces summed over sensor zone + CONTACT: contacts which occurred during the simulation ACCELEROMETER: accelerometer FORCE: force TORQUE: torque @@ -381,6 +398,7 @@ class SensorType(enum.IntEnum): TENDONLIMITFRC: tendon limit force FRAMELINACC: 3D linear acceleration FRAMEANGACC: 3D angular acceleration + TACTILE: tactile sensor """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER @@ -414,6 +432,7 @@ class SensorType(enum.IntEnum): SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM TOUCH = mujoco.mjtSensor.mjSENS_TOUCH + CONTACT = mujoco.mjtSensor.mjSENS_CONTACT ACCELEROMETER = mujoco.mjtSensor.mjSENS_ACCELEROMETER FORCE = mujoco.mjtSensor.mjSENS_FORCE TORQUE = mujoco.mjtSensor.mjSENS_TORQUE @@ -424,6 +443,7 @@ class SensorType(enum.IntEnum): TENDONLIMITFRC = mujoco.mjtSensor.mjSENS_TENDONLIMITFRC FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC + TACTILE = mujoco.mjtSensor.mjSENS_TACTILE class ObjType(enum.IntEnum): @@ -502,20 +522,6 @@ vec10 = vec10f vec11 = vec11f -class BroadphaseType(enum.IntEnum): - """Type of broadphase algorithm. - - Attributes: - NXN: Broad phase checking all pairs - SAP_TILE: Sweep and prune broad phase using tile sort - SAP_SEGMENTED: Sweep and prune broad phase using segment sort - """ - - NXN = 0 - SAP_TILE = 1 - SAP_SEGMENTED = 2 - - @dataclasses.dataclass class Option: """Physics options. @@ -527,9 +533,9 @@ class Option: ls_tolerance: CG/Newton linesearch tolerance gravity: gravitational acceleration magnetic: global magnetic flux - integrator: integration mode (mjtIntegrator) - cone: type of friction cone (mjtCone) - solver: solver algorithm (mjtSolver) + integrator: integration mode (IntegratorType) + cone: type of friction cone (ConeType) + solver: solver algorithm (SolverType) iterations: number of main solver iterations ls_iterations: maximum number of CG/Newton linesearch iterations disableflags: bit flags for disabling standard features @@ -542,8 +548,8 @@ class Option: has_fluid: True if wind, density, or viscosity are non-zero at put_model time density: density of medium viscosity: viscosity of medium - broadphase: broadphase type, 0: nxn, 1: sap_tile, 2: sap_segmented - broadphase_filter: broadphase filter bitflag + broadphase: broadphase type (BroadphaseType) + broadphase_filter: broadphase filter bitflag (BroadphaseFilter) graph_conditional: flag to use cuda graph conditional, should be False when JAX is used sdf_initpoints: number of starting points for gradient descent sdf_iterations: max number of iterations for gradient descent @@ -597,7 +603,7 @@ class Constraint: """Constraint data. Attributes: - type: constraint type (mjtConstraint) (nworld, njmax) + type: constraint type (ConstraintType) (nworld, njmax) id: id of object of specific type (nworld, njmax) J: constraint Jacobian (nworld, njmax, nv) pos: constraint position (equality, contact) (nworld, njmax) @@ -628,8 +634,6 @@ class Constraint: prev_grad: previous grad (nworld, nv) prev_Mgrad: previous Mgrad (nworld, nv) beta: polak-ribiere beta (nworld,) - beta_num: numerator of beta (nworld,) - beta_den: denominator of beta (nworld,) done: solver done (nworld,) ls_done: linesearch done (nworld,) p0: initial point (nworld, 3) @@ -684,8 +688,6 @@ class Constraint: prev_grad: wp.array2d(dtype=float) prev_Mgrad: wp.array2d(dtype=float) beta: wp.array(dtype=float) - beta_num: wp.array(dtype=float) - beta_den: wp.array(dtype=float) done: wp.array(dtype=bool) # linesearch ls_done: wp.array(dtype=bool) @@ -752,6 +754,7 @@ class Model: nwrap: number of wrap objects in all tendon paths nsensor: number of sensors nsensordata: number of elements in sensor data vector + nsensortaxel: number of taxels in all tactile sensors nmeshvert: number of vertices for all meshes nmeshface: number of faces for all meshes nmeshgraph: number of ints in mesh auxiliary data @@ -771,8 +774,6 @@ class Model: qM_mulm_i: sparse mass matrix addressing qM_mulm_j: sparse mass matrix addressing qM_madr_ij: sparse mass matrix addressing - qLD_update_tree: dof tree ordering for qLD updates - qLD_update_treeadr: index of each dof tree level M_rownnz: number of non-zeros in each row of qM (nv,) M_rowadr: index of each row in qM (nv,) M_colind: column indices of non-zeros in qM (nM,) @@ -801,7 +802,7 @@ class Model: body_contype: OR over all geom contypes (nbody,) body_conaffinity: OR over all geom conaffinities (nbody,) body_gravcomp: antigravity force, units of body weight (nworld, nbody) - jnt_type: type of joint (mjtJoint) (njnt,) + jnt_type: type of joint (JointType) (njnt,) jnt_qposadr: start addr in 'qpos' for joint's data (njnt,) jnt_dofadr: start addr in 'qvel' for joint's data (njnt,) jnt_bodyid: id of joint's body (njnt,) @@ -830,7 +831,7 @@ class Model: dof_solref: constraint solver reference: frictionloss (nworld, nv, NREF) dof_tri_row: np.tril_indices (mjm.nv)[0] dof_tri_col: np.tril_indices (mjm.nv)[1] - geom_type: geometric type (mjtGeom) (ngeom,) + geom_type: geometric type (GeomType) (ngeom,) geom_contype: geom contact type (ngeom,) geom_conaffinity: geom contact affinity (ngeom,) geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) @@ -856,11 +857,11 @@ class Model: hfield_ncol: number of columns in grid (nhfield,) hfield_size: (x, y, z_top, z_bottom) (nhfield, 4) hfield_data: elevation data (nhfielddata,) - site_type: geom type for rendering (mjtGeom) (nsite,) + site_type: geom type for rendering (GeomType) (nsite,) site_bodyid: id of site's body (nsite,) site_pos: local position offset rel. to body (nworld, nsite, 3) site_quat: local orientation offset rel. to body (nworld, nsite, 4) - cam_mode: camera tracking mode (mjtCamLight) (ncam,) + cam_mode: camera tracking mode (CamLightType) (ncam,) cam_bodyid: id of camera's body (ncam,) cam_targetbodyid: id of targeted body; -1: none (ncam,) cam_pos: position rel. to body frame (nworld, ncam, 3) @@ -872,7 +873,7 @@ class Model: cam_resolution: resolution: pixels [width, height] (ncam, 2) cam_sensorsize: sensor size: length [width, height] (ncam, 2) cam_intrinsic: [focal length; principal point] (ncam, 4) - light_mode: light tracking mode (mjtCamLight) (nlight,) + light_mode: light tracking mode (CamLightType) (nlight,) light_bodyid: id of light's body (nlight,) light_targetbodyid: id of targeted body; -1: none (nlight,) light_pos: position rel. to body frame (nworld, nlight, 3) @@ -883,10 +884,14 @@ class Model: mesh_vertadr: first vertex address (nmesh,) mesh_vertnum: number of vertices (nmesh,) mesh_vert: vertex positions for all meshes (nmeshvert, 3) + mesh_normal: normals for all meshes (nmeshnormal, 3) mesh_faceadr: first face address (nmesh,) mesh_face: face indices for all meshes (nface, 3) + mesh_normaladr: first normal address (nmesh,) + mesh_normal: normals for all meshes (nmeshnormal x 3) mesh_graphadr: graph data address; -1: no graph (nmesh,) mesh_graph: convex graph data (nmeshgraph,) + mesh_quat: rotation applied to asset vertices (nmesh, 4) mesh_polynum: number of polygons per mesh (nmesh,) mesh_polyadr: first polygon address per mesh (nmesh,) mesh_polynormal: all polygon normals (nmeshpoly, 3) @@ -896,10 +901,10 @@ class Model: mesh_polymapadr: first polygon address per vertex (nmeshvert,) mesh_polymapnum: number of polygons per vertex (nmeshvert,) mesh_polymap: vertex to polygon map (nmeshpolymap,) - eq_type: constraint type (mjtEq) (neq,) + eq_type: constraint type (EqType) (neq,) eq_obj1id: id of object 1 (neq,) eq_obj2id: id of object 2 (neq,) - eq_objtype: type of both objects (mjtObj) (neq,) + eq_objtype: type of both objects (ObjType) (neq,) eq_active0: initial enable/disable constraint state (neq,) eq_solref: constraint solver reference (nworld, neq, mjNREF) eq_solimp: constraint solver impedance (nworld, neq, mjNIMP) @@ -911,10 +916,10 @@ class Model: actuator_moment_tiles_nv: tiling configuration actuator_moment_tiles_nu: tiling configuration actuator_affine_bias_gain: affine bias/gain present - actuator_trntype: transmission type (mjtTrn) (nu,) - actuator_dyntype: dynamics type (mjtDyn) (nu,) - actuator_gaintype: gain type (mjtGain) (nu,) - actuator_biastype: bias type (mjtBias) (nu,) + actuator_trntype: transmission type (TrnType) (nu,) + actuator_dyntype: dynamics type (DynType) (nu,) + actuator_gaintype: gain type (GainType) (nu,) + actuator_biastype: bias type (BiasType) (nu,) actuator_trnid: transmission id: joint, tendon, site (nu, 2) actuator_actadr: first activation address; -1: stateless (nu,) actuator_actnum: number of activation variables (nu,) @@ -971,7 +976,7 @@ class Model: tendon_invweight0: inv. weight in qpos0 (nworld, ntendon) wrap_objid: object id: geom, site, joint (nwrap,) wrap_prm: divisor, joint coef, or site id (nwrap,) - wrap_type: wrap object type (mjtWrap) (nwrap,) + wrap_type: wrap object type (WrapType) (nwrap,) tendon_jnt_adr: joint tendon address (<=nwrap,) tendon_site_pair_adr: site pair tendon address (<=nwrap,) tendon_geom_adr: geom tendon address (<=nwrap,) @@ -982,12 +987,13 @@ class Model: wrap_site_pair_adr: first address for site wrap pair (<=nwrap,) wrap_geom_adr: addresses for geom tendon wrap object (<=nwrap,) wrap_pulley_scale: pulley scaling (nwrap,) - sensor_type: sensor type (mjtSensor) (nsensor,) - sensor_datatype: numeric data type (mjtDataType) (nsensor,) - sensor_objtype: type of sensorized object (mjtObj) (nsensor,) + sensor_type: sensor type (SensorType) (nsensor,) + sensor_datatype: numeric data type (DataType) (nsensor,) + sensor_objtype: type of sensorized object (ObjType) (nsensor,) sensor_objid: id of sensorized object (nsensor,) - sensor_reftype: type of reference frame (mjtObj) (nsensor,) + sensor_reftype: type of reference frame (ObjType) (nsensor,) sensor_refid: id of reference frame; -1: global frame (nsensor,) + sensor_intprm: sensor parameters (nsensor, mjNSENS) sensor_dim: number of scalar outputs (nsensor,) sensor_adr: address in sensor array (nsensor,) sensor_cutoff: cutoff for real and positive; 0: ignore (nsensor,) @@ -1007,6 +1013,7 @@ class Model: sensor_e_kinetic: evaluate energy_vel sensor_tendonactfrc_adr: address for tendonactfrc sensor (<=nsensor,) sensor_subtree_vel: evaluate subtree_vel + sensor_contact_adr: addresses for contact sensors sensor_rne_postconstraint: evaluate rne_postconstraint sensor_rangefinder_bodyid: bodyid for rangefinder (nrangefinder,) plugin: globally registered plugin slot number (nplugin,) @@ -1048,6 +1055,7 @@ class Model: nwrap: int nsensor: int nsensordata: int + nsensortaxel: int nmeshvert: int nmeshface: int nmeshgraph: int @@ -1193,10 +1201,13 @@ class Model: mesh_vertadr: wp.array(dtype=int) mesh_vertnum: wp.array(dtype=int) mesh_vert: wp.array(dtype=wp.vec3) + mesh_normaladr: wp.array(dtype=int) + mesh_normal: wp.array(dtype=wp.vec3) mesh_faceadr: wp.array(dtype=int) mesh_face: wp.array(dtype=wp.vec3i) mesh_graphadr: wp.array(dtype=int) mesh_graph: wp.array(dtype=int) + mesh_quat: wp.array(dtype=wp.quat) mesh_polynum: wp.array(dtype=int) mesh_polyadr: wp.array(dtype=int) mesh_polynormal: wp.array(dtype=wp.vec3) @@ -1295,6 +1306,7 @@ class Model: sensor_objid: wp.array(dtype=int) sensor_reftype: wp.array(dtype=int) sensor_refid: wp.array(dtype=int) + sensor_intprm: wp.array2d(dtype=int) sensor_dim: wp.array(dtype=int) sensor_adr: wp.array(dtype=int) sensor_cutoff: wp.array(dtype=float) @@ -1311,6 +1323,7 @@ class Model: sensor_e_kinetic: bool # warp only sensor_tendonactfrc_adr: wp.array(dtype=int) # warp only sensor_subtree_vel: bool # warp only + sensor_contact_adr: wp.array(dtype=int) # warp only sensor_rne_postconstraint: bool # warp only sensor_rangefinder_bodyid: wp.array(dtype=int) # warp only plugin: wp.array(dtype=int) @@ -1323,6 +1336,8 @@ class Model: block_dim: BlockDim # warp only geom_pair_type_count: tuple[int, ...] # warp only has_sdf_geom: bool # warp only + taxel_vertadr: wp.array(dtype=int) # warp only + taxel_sensorid: wp.array(dtype=int) # warp only @dataclasses.dataclass @@ -1365,18 +1380,19 @@ class Data: Attributes: nworld: number of worlds nconmax: maximum number of contacts - njmax: maximum number of constraints + njmax: maximum number of constraints per world solver_niter: number of solver iterations (nworld,) ncon: number of detected contacts + ncon_world: number of detected contacts per world (nworld,) ncon_hfield: number of contacts per geom pair with hfield (nworld, nhfieldgeompair) - ne: number of equality constraints - ne_connect: number of equality connect constraints - ne_weld: number of equality weld constraints - ne_jnt: number of equality joint constraints - ne_ten: number of equality tendon constraints - nf: number of friction constraints - nl: number of limit constraints - nefc: number of constraints (1,) + ne: number of equality constraints (nworld,) + ne_connect: number of equality connect constraints (nworld,) + ne_weld: number of equality weld constraints (nworld,) + ne_jnt: number of equality joint constraints (nworld,) + ne_ten: number of equality tendon constraints (nworld,) + nf: number of friction constraints (nworld,) + nl: number of limit constraints (nworld,) + nefc: number of constraints (nworld,) nsolving: number of unconverged worlds (1,) time: simulation time (nworld,) energy: potential, kinetic energy (nworld, 2) @@ -1498,6 +1514,10 @@ class Data: sensor_rangefinder_vec: directions for rangefinder (nworld, nrangefinder, 3) sensor_rangefinder_dist: distances for rangefinder (nworld, nrangefinder) sensor_rangefinder_geomid: geomids for rangefinder (nworld, nrangefinder) + sensor_contact_nmatch: match count for each world-sensor (nworld, <=nsensor) + sensor_contact_matchid: id for matching contact (nworld, <=nsensor, MJ_MAXCONPAIR) + sensor_contact_criteria: critera for reduction (nworld, <=nsensor, MJ_MAXCONPAIR) + sensor_contact_direction: direction of contact (nworld, <=nsensor, MJ_MAXCONPAIR) ray_bodyexclude: id of body to exclude from ray computation ray_dist: ray distance to nearest geom (nworld, 1) ray_geomid: id of geom that intersects with ray (nworld, 1) @@ -1510,6 +1530,7 @@ class Data: njmax: int # warp only solver_niter: wp.array(dtype=int) ncon: wp.array(dtype=int) + ncon_world: wp.array(dtype=int) # warp only ncon_hfield: wp.array2d(dtype=int) # warp only ne: wp.array(dtype=int) ne_connect: wp.array(dtype=int) # warp only @@ -1653,6 +1674,10 @@ class Data: sensor_rangefinder_vec: wp.array2d(dtype=wp.vec3) # warp only sensor_rangefinder_dist: wp.array2d(dtype=float) # warp only sensor_rangefinder_geomid: wp.array2d(dtype=int) # warp only + sensor_contact_nmatch: wp.array2d(dtype=int) # warp only + sensor_contact_matchid: wp.array3d(dtype=int) # warp only + sensor_contact_criteria: wp.array3d(dtype=float) # warp only + sensor_contact_direction: wp.array3d(dtype=float) # warp only # ray ray_bodyexclude: wp.array(dtype=int) # warp only diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py index 572d2312..aed397fd 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py @@ -14,6 +14,7 @@ # ============================================================================== import functools +import inspect from typing import Callable, Optional import warp as wp @@ -97,6 +98,11 @@ def event_scope(fn, name: str = ""): global _STACK if _STACK is None: return fn(*args, **kwargs) + + for frame_info in inspect.stack(): + if frame_info.function in ("capture_while", "capture_if"): + return fn(*args, **kwargs) + # push into next level of stack saved_stack, _STACK = _STACK, {} beg = wp.Event(enable_timing=True) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/gear.py b/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/gear.py new file mode 100644 index 00000000..c505f264 --- /dev/null +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/gear.py @@ -0,0 +1,144 @@ +import warp as wp + + +@wp.func +def Subtraction(a: float, b: float) -> float: + return wp.max(a, -b) + + +@wp.func +def Intersection(a: float, b: float) -> float: + return wp.max(a, b) + + +@wp.func +def circle(rho: float, r: float) -> float: + return rho - r + + +@wp.func +def smoothUnion(a: float, b: float, k: float) -> float: + h = wp.min(wp.max(0.5 + 0.5 * (b - a) / k, 0.0), 1.0) + return b * (1.0 - h) + a * h - k * h * (1.0 - h) + + +@wp.func +def smoothIntersection(a: float, b: float, k: float) -> float: + return Subtraction(Intersection(a, b), smoothUnion(Subtraction(a, b), Subtraction(b, a), k)) + + +@wp.func +def extrusion(p: wp.vec3, sdf_2d: float, h: float) -> float: + w = wp.vec2() + w[0] = sdf_2d + w[1] = wp.abs(p[2]) - h + w_abs = wp.vec2() + w_abs[0] = wp.max(w[0], 0.0) + w_abs[1] = wp.max(w[1], 0.0) + return wp.min(wp.max(w[0], w[1]), 0.0) + wp.sqrt(w_abs[0] * w_abs[0] + w_abs[1] * w_abs[1]) + + +@wp.func +def mod(x: float, y: float) -> float: + return x - y * wp.floor(x / y) + + +@wp.func +def distance2D(p: wp.vec3, attributes: wp.vec3) -> float: + # see https://www.shadertoy.com/view/3lG3WR + D = 2.8 + N = 25.0 + psi = 3.096e-5 * N * N - 6.557e-3 * N + 0.551 # pressure angle + alpha = 0.0 + innerdiameter = -1.0 + + R = D / 2.0 + rho = wp.sqrt(p[0] * p[0] + p[1] * p[1]) + Pd = N / D # Diametral Pitch: teeth per unit length of diameter + P = wp.PI / Pd # Circular Pitch + a = 1.0 / Pd # Addendum: radial length of a tooth from the pitch + # circle to the tip of the tooth. + + Do = D + 2.0 * a # Outside Diameter + Ro = Do / 2.0 + + h = 2.2 / Pd + + innerR = Ro - h - 0.14 * D + if innerdiameter >= 0.0: + innerR = innerdiameter / 2.0 + + # Early exit + if innerR - rho > 0.0: + return innerR - rho + + # Early exit + if Ro - rho < -0.2: + return rho - Ro + + Db = D * wp.cos(psi) # Base Diameter + Rb = Db / 2.0 + + fi = wp.atan2(p[1], p[0]) + alpha + alphaStride = P / R + + invAlpha = wp.acos(Rb / R) + invPhi = wp.tan(invAlpha) - invAlpha + + shift = alphaStride / 2.0 - 2.0 * invPhi + + fia = mod(fi + shift / 2.0, alphaStride) - shift / 2.0 + fib = mod(-fi - shift + shift / 2.0, alphaStride) - shift / 2.0 + + dista = -1.0e6 + distb = -1.0e6 + + if Rb < rho: + acos_rbRho = wp.acos(Rb / rho) + + thetaa = fia + acos_rbRho + thetab = fib + acos_rbRho + + ta = wp.sqrt(rho * rho - Rb * Rb) + + # https://math.stackexchange.com/questions/1266689/distance-from-a-point-to-the-involute-of-a-circle + dista = ta - Rb * thetaa + distb = ta - Rb * thetab + + gearOuter = circle(rho, Ro) + gearLowBase = circle(rho, Ro - h) + crownBase = circle(rho, innerR) + cogs = Intersection(dista, distb) + baseWalls = Intersection(fia - (alphaStride - shift), fib - (alphaStride - shift)) + + cogs = Intersection(baseWalls, cogs) + cogs = smoothIntersection(gearOuter, cogs, 0.0035 * D) + cogs = smoothUnion(gearLowBase, cogs, Rb - Ro + h) + cogs = Subtraction(cogs, crownBase) + + return cogs + + +@wp.func +def gear(p: wp.vec3, attr: wp.vec3) -> float: + thickness = 0.2 + return extrusion(p, distance2D(p, attr), thickness / 2.0) + + +@wp.func +def gear_sdf_grad(p: wp.vec3, attr: wp.vec3) -> wp.vec3: + grad = wp.vec3() + eps = 1e-6 + f_original = gear(p, attr) + x_plus = wp.vec3(p[0] + eps, p[1], p[2]) + f_plus = gear(x_plus, attr) + grad[0] = (f_plus - f_original) / eps + + x_plus = wp.vec3(p[0], p[1] + eps, p[2]) + f_plus = gear(x_plus, attr) + grad[1] = (f_plus - f_original) / eps + + x_plus = wp.vec3(p[0], p[1], p[2] + eps) + f_plus = gear(x_plus, attr) + grad[2] = (f_plus - f_original) / eps + return grad diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/tactile.xml b/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/tactile.xml new file mode 100644 index 00000000..470611a9 --- /dev/null +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/tactile.xml @@ -0,0 +1,66 @@ + + + diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/utils.py b/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/utils.py index 76a12f89..5526dff9 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/utils.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/test_data/collision_sdf/utils.py @@ -22,6 +22,8 @@ import warp as wp from .bolt import bolt from .bolt import bolt_sdf_grad +from .gear import gear +from .gear import gear_sdf_grad from .nut import nut from .nut import nut_sdf_grad @@ -31,21 +33,25 @@ class SDFType(enum.Enum): NUT = "NUT" BOLT = "BOLT" + GEAR = "GEAR" -def register_sdf_plugins(collision_sdf) -> Dict[str, int]: +def register_sdf_plugins(mjwarp) -> Dict[str, int]: xml = """ + + + """ @@ -62,6 +68,8 @@ def register_sdf_plugins(collision_sdf) -> Dict[str, int]: sdf_types[SDFType.NUT.value] = int(m.plugin[i]) elif name == "bg": sdf_types[SDFType.BOLT.value] = int(m.plugin[i]) + elif name == "gg": + sdf_types[SDFType.GEAR.value] = int(m.plugin[i]) @wp.func def user_sdf(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float: @@ -70,6 +78,8 @@ def register_sdf_plugins(collision_sdf) -> Dict[str, int]: result = nut(p, attr) elif sdf_type == wp.static(sdf_types[SDFType.BOLT.value]): result = bolt(p, attr) + elif sdf_type == wp.static(sdf_types[SDFType.GEAR.value]): + result = gear(p, attr) return result @wp.func @@ -78,9 +88,11 @@ def register_sdf_plugins(collision_sdf) -> Dict[str, int]: return nut_sdf_grad(p, attr) elif sdf_type == wp.static(sdf_types[SDFType.BOLT.value]): return bolt_sdf_grad(p, attr) + elif sdf_type == wp.static(sdf_types[SDFType.GEAR.value]): + return gear_sdf_grad(p, attr) return wp.vec3() - collision_sdf.user_sdf = user_sdf - collision_sdf.user_sdf_grad = user_sdf_grad + mjwarp._src.collision_sdf.user_sdf = user_sdf + mjwarp._src.collision_sdf.user_sdf_grad = user_sdf_grad return sdf_types diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 0b86774e..ecc7c0d4 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -42,7 +42,6 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) - @ffi.format_args_for_warp def _collision_shim( # Model @@ -489,8 +488,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def collision(m: types.Model, d: types.Data): return _collision_jax_impl(m, d) - - @collision.def_vmap @ffi.marshal_custom_vmap def collision_vmap(unused_axis_size, is_batched, m, d): diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 8238c24d..aa54840f 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -195,6 +195,8 @@ def _forward_shim( mesh_faceadr: wp.array(dtype=int), mesh_graph: wp.array(dtype=int), mesh_graphadr: wp.array(dtype=int), + mesh_normal: wp.array(dtype=wp.vec3), + mesh_normaladr: wp.array(dtype=int), mesh_polyadr: wp.array(dtype=int), mesh_polymap: wp.array(dtype=int), mesh_polymapadr: wp.array(dtype=int), @@ -204,6 +206,7 @@ def _forward_shim( mesh_polyvert: wp.array(dtype=int), mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertnum: wp.array(dtype=int), + mesh_quat: wp.array(dtype=wp.quat), mesh_vert: wp.array(dtype=wp.vec3), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), @@ -224,6 +227,8 @@ def _forward_shim( nlsp: int, nmeshface: int, nmocap: int, + nsensordata: int, + nsensortaxel: int, nsite: int, ntendon: int, nu: int, @@ -252,10 +257,13 @@ def _forward_shim( rangefinder_sensor_adr: wp.array(dtype=int), sensor_acc_adr: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), + sensor_contact_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), sensor_datatype: wp.array(dtype=int), + sensor_dim: wp.array(dtype=int), sensor_e_kinetic: bool, sensor_e_potential: bool, + sensor_intprm: wp.array2d(dtype=int), sensor_limitfrc_adr: wp.array(dtype=int), sensor_limitpos_adr: wp.array(dtype=int), sensor_limitvel_adr: wp.array(dtype=int), @@ -278,6 +286,8 @@ def _forward_shim( site_size: wp.array(dtype=wp.vec3), site_type: wp.array(dtype=int), subtree_mass: wp.array2d(dtype=float), + taxel_sensorid: wp.array(dtype=int), + taxel_vertadr: wp.array(dtype=int), tendon_actfrclimited: wp.array(dtype=bool), tendon_actfrcrange: wp.array2d(dtype=wp.vec2), tendon_adr: wp.array(dtype=int), @@ -418,6 +428,10 @@ def _forward_shim( sap_range: wp.array2d(dtype=int), sap_segment_index: wp.array2d(dtype=int), sap_sort_index: wp.array3d(dtype=int), + sensor_contact_criteria: wp.array3d(dtype=float), + sensor_contact_direction: wp.array3d(dtype=float), + sensor_contact_matchid: wp.array3d(dtype=int), + sensor_contact_nmatch: wp.array2d(dtype=int), sensor_rangefinder_dist: wp.array2d(dtype=float), sensor_rangefinder_geomid: wp.array2d(dtype=int), sensor_rangefinder_pnt: wp.array2d(dtype=wp.vec3), @@ -471,8 +485,6 @@ def _forward_shim( efc__alpha: wp.array(dtype=float), efc__aref: wp.array2d(dtype=float), efc__beta: wp.array(dtype=float), - efc__beta_den: wp.array(dtype=float), - efc__beta_num: wp.array(dtype=float), efc__cholesky_L_tmp: wp.array3d(dtype=float), efc__cholesky_y_tmp: wp.array2d(dtype=float), efc__condim: wp.array2d(dtype=int), @@ -669,6 +681,8 @@ def _forward_shim( _m.mesh_faceadr = mesh_faceadr _m.mesh_graph = mesh_graph _m.mesh_graphadr = mesh_graphadr + _m.mesh_normal = mesh_normal + _m.mesh_normaladr = mesh_normaladr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr @@ -678,6 +692,7 @@ def _forward_shim( _m.mesh_polyvert = mesh_polyvert _m.mesh_polyvertadr = mesh_polyvertadr _m.mesh_polyvertnum = mesh_polyvertnum + _m.mesh_quat = mesh_quat _m.mesh_vert = mesh_vert _m.mesh_vertadr = mesh_vertadr _m.mesh_vertnum = mesh_vertnum @@ -698,6 +713,8 @@ def _forward_shim( _m.nlsp = nlsp _m.nmeshface = nmeshface _m.nmocap = nmocap + _m.nsensordata = nsensordata + _m.nsensortaxel = nsensortaxel _m.nsite = nsite _m.ntendon = ntendon _m.nu = nu @@ -752,10 +769,13 @@ def _forward_shim( _m.rangefinder_sensor_adr = rangefinder_sensor_adr _m.sensor_acc_adr = sensor_acc_adr _m.sensor_adr = sensor_adr + _m.sensor_contact_adr = sensor_contact_adr _m.sensor_cutoff = sensor_cutoff _m.sensor_datatype = sensor_datatype + _m.sensor_dim = sensor_dim _m.sensor_e_kinetic = sensor_e_kinetic _m.sensor_e_potential = sensor_e_potential + _m.sensor_intprm = sensor_intprm _m.sensor_limitfrc_adr = sensor_limitfrc_adr _m.sensor_limitpos_adr = sensor_limitpos_adr _m.sensor_limitvel_adr = sensor_limitvel_adr @@ -779,6 +799,8 @@ def _forward_shim( _m.site_type = site_type _m.stat.meaninertia = stat__meaninertia _m.subtree_mass = subtree_mass + _m.taxel_sensorid = taxel_sensorid + _m.taxel_vertadr = taxel_vertadr _m.tendon_actfrclimited = tendon_actfrclimited _m.tendon_actfrcrange = tendon_actfrcrange _m.tendon_adr = tendon_adr @@ -850,8 +872,6 @@ def _forward_shim( _d.efc.alpha = efc__alpha _d.efc.aref = efc__aref _d.efc.beta = efc__beta - _d.efc.beta_den = efc__beta_den - _d.efc.beta_num = efc__beta_num _d.efc.cholesky_L_tmp = efc__cholesky_L_tmp _d.efc.cholesky_y_tmp = efc__cholesky_y_tmp _d.efc.condim = efc__condim @@ -957,6 +977,10 @@ def _forward_shim( _d.sap_range = sap_range _d.sap_segment_index = sap_segment_index _d.sap_sort_index = sap_sort_index + _d.sensor_contact_criteria = sensor_contact_criteria + _d.sensor_contact_direction = sensor_contact_direction + _d.sensor_contact_matchid = sensor_contact_matchid + _d.sensor_contact_nmatch = sensor_contact_nmatch _d.sensor_rangefinder_dist = sensor_rangefinder_dist _d.sensor_rangefinder_geomid = sensor_rangefinder_geomid _d.sensor_rangefinder_pnt = sensor_rangefinder_pnt @@ -1077,6 +1101,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'sap_range': d._impl.sap_range.shape, 'sap_segment_index': d._impl.sap_segment_index.shape, 'sap_sort_index': d._impl.sap_sort_index.shape, + 'sensor_contact_criteria': d._impl.sensor_contact_criteria.shape, + 'sensor_contact_direction': d._impl.sensor_contact_direction.shape, + 'sensor_contact_matchid': d._impl.sensor_contact_matchid.shape, + 'sensor_contact_nmatch': d._impl.sensor_contact_nmatch.shape, 'sensor_rangefinder_dist': d._impl.sensor_rangefinder_dist.shape, 'sensor_rangefinder_geomid': d._impl.sensor_rangefinder_geomid.shape, 'sensor_rangefinder_pnt': d._impl.sensor_rangefinder_pnt.shape, @@ -1130,8 +1158,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__alpha': d._impl.efc__alpha.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__beta': d._impl.efc__beta.shape, - 'efc__beta_den': d._impl.efc__beta_den.shape, - 'efc__beta_num': d._impl.efc__beta_num.shape, 'efc__cholesky_L_tmp': d._impl.efc__cholesky_L_tmp.shape, 'efc__cholesky_y_tmp': d._impl.efc__cholesky_y_tmp.shape, 'efc__condim': d._impl.efc__condim.shape, @@ -1178,7 +1204,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _forward_shim, - num_outputs=180, + num_outputs=182, output_dims=output_dims, vmap_method=None, graph_compatible=True, @@ -1265,6 +1291,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'sap_range', 'sap_segment_index', 'sap_sort_index', + 'sensor_contact_criteria', + 'sensor_contact_direction', + 'sensor_contact_matchid', + 'sensor_contact_nmatch', 'sensor_rangefinder_dist', 'sensor_rangefinder_geomid', 'sensor_rangefinder_pnt', @@ -1318,8 +1348,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__alpha', 'efc__aref', 'efc__beta', - 'efc__beta_den', - 'efc__beta_num', 'efc__cholesky_L_tmp', 'efc__cholesky_y_tmp', 'efc__condim', @@ -1515,6 +1543,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.mesh_faceadr, m.mesh_graph, m.mesh_graphadr, + m.mesh_normal, + m.mesh_normaladr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, @@ -1524,6 +1554,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.mesh_polyvert, m._impl.mesh_polyvertadr, m._impl.mesh_polyvertnum, + m.mesh_quat, m.mesh_vert, m.mesh_vertadr, m.mesh_vertnum, @@ -1544,6 +1575,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.nlsp, m.nmeshface, m.nmocap, + m.nsensordata, + m._impl.nsensortaxel, m.nsite, m.ntendon, m.nu, @@ -1572,10 +1605,13 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.rangefinder_sensor_adr, m._impl.sensor_acc_adr, m.sensor_adr, + m._impl.sensor_contact_adr, m.sensor_cutoff, m.sensor_datatype, + m.sensor_dim, m._impl.sensor_e_kinetic, m._impl.sensor_e_potential, + m.sensor_intprm, m._impl.sensor_limitfrc_adr, m._impl.sensor_limitpos_adr, m._impl.sensor_limitvel_adr, @@ -1598,6 +1634,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.site_size, m.site_type, m._impl.subtree_mass, + m._impl.taxel_sensorid, + m._impl.taxel_vertadr, m.tendon_actfrclimited, m.tendon_actfrcrange, m.tendon_adr, @@ -1737,6 +1775,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.sap_range, d._impl.sap_segment_index, d._impl.sap_sort_index, + d._impl.sensor_contact_criteria, + d._impl.sensor_contact_direction, + d._impl.sensor_contact_matchid, + d._impl.sensor_contact_nmatch, d._impl.sensor_rangefinder_dist, d._impl.sensor_rangefinder_geomid, d._impl.sensor_rangefinder_pnt, @@ -1790,8 +1832,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.efc__alpha, d._impl.efc__aref, d._impl.efc__beta, - d._impl.efc__beta_den, - d._impl.efc__beta_num, d._impl.efc__cholesky_L_tmp, d._impl.efc__cholesky_y_tmp, d._impl.efc__condim, @@ -1919,104 +1959,106 @@ def _forward_jax_impl(m: types.Model, d: types.Data): '_impl.sap_range': out[79], '_impl.sap_segment_index': out[80], '_impl.sap_sort_index': out[81], - '_impl.sensor_rangefinder_dist': out[82], - '_impl.sensor_rangefinder_geomid': out[83], - '_impl.sensor_rangefinder_pnt': out[84], - '_impl.sensor_rangefinder_vec': out[85], - 'sensordata': out[86], - 'site_xmat': out[87], - 'site_xpos': out[88], - '_impl.solver_niter': out[89], - '_impl.subtree_angmom': out[90], - '_impl.subtree_bodyvel': out[91], - 'subtree_com': out[92], - '_impl.subtree_linvel': out[93], - '_impl.ten_J': out[94], - '_impl.ten_Jdot': out[95], - '_impl.ten_actfrc': out[96], - '_impl.ten_bias_coef': out[97], - '_impl.ten_length': out[98], - '_impl.ten_velocity': out[99], - '_impl.ten_wrapadr': out[100], - '_impl.ten_wrapnum': out[101], - 'time': out[102], - '_impl.wrap_geom_xpos': out[103], - '_impl.wrap_obj': out[104], - '_impl.wrap_xpos': out[105], - 'xanchor': out[106], - 'xaxis': out[107], - 'xfrc_applied': out[108], - 'ximat': out[109], - 'xipos': out[110], - 'xmat': out[111], - 'xpos': out[112], - 'xquat': out[113], - '_impl.contact__dim': out[114], - '_impl.contact__dist': out[115], - '_impl.contact__efc_address': out[116], - '_impl.contact__frame': out[117], - '_impl.contact__friction': out[118], - '_impl.contact__geom': out[119], - '_impl.contact__includemargin': out[120], - '_impl.contact__pos': out[121], - '_impl.contact__solimp': out[122], - '_impl.contact__solref': out[123], - '_impl.contact__solreffriction': out[124], - '_impl.contact__worldid': out[125], - '_impl.efc__D': out[126], - '_impl.efc__J': out[127], - '_impl.efc__Jaref': out[128], - '_impl.efc__Ma': out[129], - '_impl.efc__Mgrad': out[130], - '_impl.efc__active': out[131], - '_impl.efc__alpha': out[132], - '_impl.efc__aref': out[133], - '_impl.efc__beta': out[134], - '_impl.efc__beta_den': out[135], - '_impl.efc__beta_num': out[136], - '_impl.efc__cholesky_L_tmp': out[137], - '_impl.efc__cholesky_y_tmp': out[138], - '_impl.efc__condim': out[139], - '_impl.efc__cost': out[140], - '_impl.efc__cost_candidate': out[141], - '_impl.efc__done': out[142], - '_impl.efc__force': out[143], - '_impl.efc__frictionloss': out[144], - '_impl.efc__gauss': out[145], - '_impl.efc__grad': out[146], - '_impl.efc__grad_dot': out[147], - '_impl.efc__gtol': out[148], - '_impl.efc__h': out[149], - '_impl.efc__hi': out[150], - '_impl.efc__hi_alpha': out[151], - '_impl.efc__hi_next': out[152], - '_impl.efc__hi_next_alpha': out[153], - '_impl.efc__id': out[154], - '_impl.efc__jv': out[155], - '_impl.efc__lo': out[156], - '_impl.efc__lo_alpha': out[157], - '_impl.efc__lo_next': out[158], - '_impl.efc__lo_next_alpha': out[159], - '_impl.efc__ls_done': out[160], - '_impl.efc__margin': out[161], - '_impl.efc__mid': out[162], - '_impl.efc__mid_alpha': out[163], - '_impl.efc__mv': out[164], - '_impl.efc__p0': out[165], - '_impl.efc__pos': out[166], - '_impl.efc__prev_Mgrad': out[167], - '_impl.efc__prev_cost': out[168], - '_impl.efc__prev_grad': out[169], - '_impl.efc__quad': out[170], - '_impl.efc__quad_gauss': out[171], - '_impl.efc__search': out[172], - '_impl.efc__search_dot': out[173], - '_impl.efc__type': out[174], - '_impl.efc__u': out[175], - '_impl.efc__uu': out[176], - '_impl.efc__uv': out[177], - '_impl.efc__vel': out[178], - '_impl.efc__vv': out[179], + '_impl.sensor_contact_criteria': out[82], + '_impl.sensor_contact_direction': out[83], + '_impl.sensor_contact_matchid': out[84], + '_impl.sensor_contact_nmatch': out[85], + '_impl.sensor_rangefinder_dist': out[86], + '_impl.sensor_rangefinder_geomid': out[87], + '_impl.sensor_rangefinder_pnt': out[88], + '_impl.sensor_rangefinder_vec': out[89], + 'sensordata': out[90], + 'site_xmat': out[91], + 'site_xpos': out[92], + '_impl.solver_niter': out[93], + '_impl.subtree_angmom': out[94], + '_impl.subtree_bodyvel': out[95], + 'subtree_com': out[96], + '_impl.subtree_linvel': out[97], + '_impl.ten_J': out[98], + '_impl.ten_Jdot': out[99], + '_impl.ten_actfrc': out[100], + '_impl.ten_bias_coef': out[101], + '_impl.ten_length': out[102], + '_impl.ten_velocity': out[103], + '_impl.ten_wrapadr': out[104], + '_impl.ten_wrapnum': out[105], + 'time': out[106], + '_impl.wrap_geom_xpos': out[107], + '_impl.wrap_obj': out[108], + '_impl.wrap_xpos': out[109], + 'xanchor': out[110], + 'xaxis': out[111], + 'xfrc_applied': out[112], + 'ximat': out[113], + 'xipos': out[114], + 'xmat': out[115], + 'xpos': out[116], + 'xquat': out[117], + '_impl.contact__dim': out[118], + '_impl.contact__dist': out[119], + '_impl.contact__efc_address': out[120], + '_impl.contact__frame': out[121], + '_impl.contact__friction': out[122], + '_impl.contact__geom': out[123], + '_impl.contact__includemargin': out[124], + '_impl.contact__pos': out[125], + '_impl.contact__solimp': out[126], + '_impl.contact__solref': out[127], + '_impl.contact__solreffriction': out[128], + '_impl.contact__worldid': out[129], + '_impl.efc__D': out[130], + '_impl.efc__J': out[131], + '_impl.efc__Jaref': out[132], + '_impl.efc__Ma': out[133], + '_impl.efc__Mgrad': out[134], + '_impl.efc__active': out[135], + '_impl.efc__alpha': out[136], + '_impl.efc__aref': out[137], + '_impl.efc__beta': out[138], + '_impl.efc__cholesky_L_tmp': out[139], + '_impl.efc__cholesky_y_tmp': out[140], + '_impl.efc__condim': out[141], + '_impl.efc__cost': out[142], + '_impl.efc__cost_candidate': out[143], + '_impl.efc__done': out[144], + '_impl.efc__force': out[145], + '_impl.efc__frictionloss': out[146], + '_impl.efc__gauss': out[147], + '_impl.efc__grad': out[148], + '_impl.efc__grad_dot': out[149], + '_impl.efc__gtol': out[150], + '_impl.efc__h': out[151], + '_impl.efc__hi': out[152], + '_impl.efc__hi_alpha': out[153], + '_impl.efc__hi_next': out[154], + '_impl.efc__hi_next_alpha': out[155], + '_impl.efc__id': out[156], + '_impl.efc__jv': out[157], + '_impl.efc__lo': out[158], + '_impl.efc__lo_alpha': out[159], + '_impl.efc__lo_next': out[160], + '_impl.efc__lo_next_alpha': out[161], + '_impl.efc__ls_done': out[162], + '_impl.efc__margin': out[163], + '_impl.efc__mid': out[164], + '_impl.efc__mid_alpha': out[165], + '_impl.efc__mv': out[166], + '_impl.efc__p0': out[167], + '_impl.efc__pos': out[168], + '_impl.efc__prev_Mgrad': out[169], + '_impl.efc__prev_cost': out[170], + '_impl.efc__prev_grad': out[171], + '_impl.efc__quad': out[172], + '_impl.efc__quad_gauss': out[173], + '_impl.efc__search': out[174], + '_impl.efc__search_dot': out[175], + '_impl.efc__type': out[176], + '_impl.efc__u': out[177], + '_impl.efc__uu': out[178], + '_impl.efc__uv': out[179], + '_impl.efc__vel': out[180], + '_impl.efc__vv': out[181], }) return d @@ -2025,8 +2067,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def forward(m: types.Model, d: types.Data): return _forward_jax_impl(m, d) - - @forward.def_vmap @ffi.marshal_custom_vmap def forward_vmap(unused_axis_size, is_batched, m, d): @@ -2207,6 +2247,8 @@ def _step_shim( mesh_faceadr: wp.array(dtype=int), mesh_graph: wp.array(dtype=int), mesh_graphadr: wp.array(dtype=int), + mesh_normal: wp.array(dtype=wp.vec3), + mesh_normaladr: wp.array(dtype=int), mesh_polyadr: wp.array(dtype=int), mesh_polymap: wp.array(dtype=int), mesh_polymapadr: wp.array(dtype=int), @@ -2216,6 +2258,7 @@ def _step_shim( mesh_polyvert: wp.array(dtype=int), mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertnum: wp.array(dtype=int), + mesh_quat: wp.array(dtype=wp.quat), mesh_vert: wp.array(dtype=wp.vec3), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), @@ -2236,6 +2279,8 @@ def _step_shim( nlsp: int, nmeshface: int, nmocap: int, + nsensordata: int, + nsensortaxel: int, nsite: int, ntendon: int, nu: int, @@ -2264,10 +2309,13 @@ def _step_shim( rangefinder_sensor_adr: wp.array(dtype=int), sensor_acc_adr: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), + sensor_contact_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), sensor_datatype: wp.array(dtype=int), + sensor_dim: wp.array(dtype=int), sensor_e_kinetic: bool, sensor_e_potential: bool, + sensor_intprm: wp.array2d(dtype=int), sensor_limitfrc_adr: wp.array(dtype=int), sensor_limitpos_adr: wp.array(dtype=int), sensor_limitvel_adr: wp.array(dtype=int), @@ -2290,6 +2338,8 @@ def _step_shim( site_size: wp.array(dtype=wp.vec3), site_type: wp.array(dtype=int), subtree_mass: wp.array2d(dtype=float), + taxel_sensorid: wp.array(dtype=int), + taxel_vertadr: wp.array(dtype=int), tendon_actfrclimited: wp.array(dtype=bool), tendon_actfrcrange: wp.array2d(dtype=wp.vec2), tendon_adr: wp.array(dtype=int), @@ -2443,6 +2493,10 @@ def _step_shim( sap_range: wp.array2d(dtype=int), sap_segment_index: wp.array2d(dtype=int), sap_sort_index: wp.array3d(dtype=int), + sensor_contact_criteria: wp.array3d(dtype=float), + sensor_contact_direction: wp.array3d(dtype=float), + sensor_contact_matchid: wp.array3d(dtype=int), + sensor_contact_nmatch: wp.array2d(dtype=int), sensor_rangefinder_dist: wp.array2d(dtype=float), sensor_rangefinder_geomid: wp.array2d(dtype=int), sensor_rangefinder_pnt: wp.array2d(dtype=wp.vec3), @@ -2496,8 +2550,6 @@ def _step_shim( efc__alpha: wp.array(dtype=float), efc__aref: wp.array2d(dtype=float), efc__beta: wp.array(dtype=float), - efc__beta_den: wp.array(dtype=float), - efc__beta_num: wp.array(dtype=float), efc__cholesky_L_tmp: wp.array3d(dtype=float), efc__cholesky_y_tmp: wp.array2d(dtype=float), efc__condim: wp.array2d(dtype=int), @@ -2695,6 +2747,8 @@ def _step_shim( _m.mesh_faceadr = mesh_faceadr _m.mesh_graph = mesh_graph _m.mesh_graphadr = mesh_graphadr + _m.mesh_normal = mesh_normal + _m.mesh_normaladr = mesh_normaladr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr @@ -2704,6 +2758,7 @@ def _step_shim( _m.mesh_polyvert = mesh_polyvert _m.mesh_polyvertadr = mesh_polyvertadr _m.mesh_polyvertnum = mesh_polyvertnum + _m.mesh_quat = mesh_quat _m.mesh_vert = mesh_vert _m.mesh_vertadr = mesh_vertadr _m.mesh_vertnum = mesh_vertnum @@ -2724,6 +2779,8 @@ def _step_shim( _m.nlsp = nlsp _m.nmeshface = nmeshface _m.nmocap = nmocap + _m.nsensordata = nsensordata + _m.nsensortaxel = nsensortaxel _m.nsite = nsite _m.ntendon = ntendon _m.nu = nu @@ -2779,10 +2836,13 @@ def _step_shim( _m.rangefinder_sensor_adr = rangefinder_sensor_adr _m.sensor_acc_adr = sensor_acc_adr _m.sensor_adr = sensor_adr + _m.sensor_contact_adr = sensor_contact_adr _m.sensor_cutoff = sensor_cutoff _m.sensor_datatype = sensor_datatype + _m.sensor_dim = sensor_dim _m.sensor_e_kinetic = sensor_e_kinetic _m.sensor_e_potential = sensor_e_potential + _m.sensor_intprm = sensor_intprm _m.sensor_limitfrc_adr = sensor_limitfrc_adr _m.sensor_limitpos_adr = sensor_limitpos_adr _m.sensor_limitvel_adr = sensor_limitvel_adr @@ -2806,6 +2866,8 @@ def _step_shim( _m.site_type = site_type _m.stat.meaninertia = stat__meaninertia _m.subtree_mass = subtree_mass + _m.taxel_sensorid = taxel_sensorid + _m.taxel_vertadr = taxel_vertadr _m.tendon_actfrclimited = tendon_actfrclimited _m.tendon_actfrcrange = tendon_actfrcrange _m.tendon_adr = tendon_adr @@ -2879,8 +2941,6 @@ def _step_shim( _d.efc.alpha = efc__alpha _d.efc.aref = efc__aref _d.efc.beta = efc__beta - _d.efc.beta_den = efc__beta_den - _d.efc.beta_num = efc__beta_num _d.efc.cholesky_L_tmp = efc__cholesky_L_tmp _d.efc.cholesky_y_tmp = efc__cholesky_y_tmp _d.efc.condim = efc__condim @@ -2996,6 +3056,10 @@ def _step_shim( _d.sap_range = sap_range _d.sap_segment_index = sap_segment_index _d.sap_sort_index = sap_sort_index + _d.sensor_contact_criteria = sensor_contact_criteria + _d.sensor_contact_direction = sensor_contact_direction + _d.sensor_contact_matchid = sensor_contact_matchid + _d.sensor_contact_nmatch = sensor_contact_nmatch _d.sensor_rangefinder_dist = sensor_rangefinder_dist _d.sensor_rangefinder_geomid = sensor_rangefinder_geomid _d.sensor_rangefinder_pnt = sensor_rangefinder_pnt @@ -3128,6 +3192,10 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'sap_range': d._impl.sap_range.shape, 'sap_segment_index': d._impl.sap_segment_index.shape, 'sap_sort_index': d._impl.sap_sort_index.shape, + 'sensor_contact_criteria': d._impl.sensor_contact_criteria.shape, + 'sensor_contact_direction': d._impl.sensor_contact_direction.shape, + 'sensor_contact_matchid': d._impl.sensor_contact_matchid.shape, + 'sensor_contact_nmatch': d._impl.sensor_contact_nmatch.shape, 'sensor_rangefinder_dist': d._impl.sensor_rangefinder_dist.shape, 'sensor_rangefinder_geomid': d._impl.sensor_rangefinder_geomid.shape, 'sensor_rangefinder_pnt': d._impl.sensor_rangefinder_pnt.shape, @@ -3181,8 +3249,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__alpha': d._impl.efc__alpha.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__beta': d._impl.efc__beta.shape, - 'efc__beta_den': d._impl.efc__beta_den.shape, - 'efc__beta_num': d._impl.efc__beta_num.shape, 'efc__cholesky_L_tmp': d._impl.efc__cholesky_L_tmp.shape, 'efc__cholesky_y_tmp': d._impl.efc__cholesky_y_tmp.shape, 'efc__condim': d._impl.efc__condim.shape, @@ -3229,7 +3295,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _step_shim, - num_outputs=192, + num_outputs=194, output_dims=output_dims, vmap_method=None, graph_compatible=True, @@ -3328,6 +3394,10 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'sap_range', 'sap_segment_index', 'sap_sort_index', + 'sensor_contact_criteria', + 'sensor_contact_direction', + 'sensor_contact_matchid', + 'sensor_contact_nmatch', 'sensor_rangefinder_dist', 'sensor_rangefinder_geomid', 'sensor_rangefinder_pnt', @@ -3381,8 +3451,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__alpha', 'efc__aref', 'efc__beta', - 'efc__beta_den', - 'efc__beta_num', 'efc__cholesky_L_tmp', 'efc__cholesky_y_tmp', 'efc__condim', @@ -3579,6 +3647,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.mesh_faceadr, m.mesh_graph, m.mesh_graphadr, + m.mesh_normal, + m.mesh_normaladr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, @@ -3588,6 +3658,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.mesh_polyvert, m._impl.mesh_polyvertadr, m._impl.mesh_polyvertnum, + m.mesh_quat, m.mesh_vert, m.mesh_vertadr, m.mesh_vertnum, @@ -3608,6 +3679,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.nlsp, m.nmeshface, m.nmocap, + m.nsensordata, + m._impl.nsensortaxel, m.nsite, m.ntendon, m.nu, @@ -3636,10 +3709,13 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.rangefinder_sensor_adr, m._impl.sensor_acc_adr, m.sensor_adr, + m._impl.sensor_contact_adr, m.sensor_cutoff, m.sensor_datatype, + m.sensor_dim, m._impl.sensor_e_kinetic, m._impl.sensor_e_potential, + m.sensor_intprm, m._impl.sensor_limitfrc_adr, m._impl.sensor_limitpos_adr, m._impl.sensor_limitvel_adr, @@ -3662,6 +3738,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.site_size, m.site_type, m._impl.subtree_mass, + m._impl.taxel_sensorid, + m._impl.taxel_vertadr, m.tendon_actfrclimited, m.tendon_actfrcrange, m.tendon_adr, @@ -3814,6 +3892,10 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.sap_range, d._impl.sap_segment_index, d._impl.sap_sort_index, + d._impl.sensor_contact_criteria, + d._impl.sensor_contact_direction, + d._impl.sensor_contact_matchid, + d._impl.sensor_contact_nmatch, d._impl.sensor_rangefinder_dist, d._impl.sensor_rangefinder_geomid, d._impl.sensor_rangefinder_pnt, @@ -3867,8 +3949,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.efc__alpha, d._impl.efc__aref, d._impl.efc__beta, - d._impl.efc__beta_den, - d._impl.efc__beta_num, d._impl.efc__cholesky_L_tmp, d._impl.efc__cholesky_y_tmp, d._impl.efc__condim, @@ -4008,104 +4088,106 @@ def _step_jax_impl(m: types.Model, d: types.Data): '_impl.sap_range': out[91], '_impl.sap_segment_index': out[92], '_impl.sap_sort_index': out[93], - '_impl.sensor_rangefinder_dist': out[94], - '_impl.sensor_rangefinder_geomid': out[95], - '_impl.sensor_rangefinder_pnt': out[96], - '_impl.sensor_rangefinder_vec': out[97], - 'sensordata': out[98], - 'site_xmat': out[99], - 'site_xpos': out[100], - '_impl.solver_niter': out[101], - '_impl.subtree_angmom': out[102], - '_impl.subtree_bodyvel': out[103], - 'subtree_com': out[104], - '_impl.subtree_linvel': out[105], - '_impl.ten_J': out[106], - '_impl.ten_Jdot': out[107], - '_impl.ten_actfrc': out[108], - '_impl.ten_bias_coef': out[109], - '_impl.ten_length': out[110], - '_impl.ten_velocity': out[111], - '_impl.ten_wrapadr': out[112], - '_impl.ten_wrapnum': out[113], - 'time': out[114], - '_impl.wrap_geom_xpos': out[115], - '_impl.wrap_obj': out[116], - '_impl.wrap_xpos': out[117], - 'xanchor': out[118], - 'xaxis': out[119], - 'xfrc_applied': out[120], - 'ximat': out[121], - 'xipos': out[122], - 'xmat': out[123], - 'xpos': out[124], - 'xquat': out[125], - '_impl.contact__dim': out[126], - '_impl.contact__dist': out[127], - '_impl.contact__efc_address': out[128], - '_impl.contact__frame': out[129], - '_impl.contact__friction': out[130], - '_impl.contact__geom': out[131], - '_impl.contact__includemargin': out[132], - '_impl.contact__pos': out[133], - '_impl.contact__solimp': out[134], - '_impl.contact__solref': out[135], - '_impl.contact__solreffriction': out[136], - '_impl.contact__worldid': out[137], - '_impl.efc__D': out[138], - '_impl.efc__J': out[139], - '_impl.efc__Jaref': out[140], - '_impl.efc__Ma': out[141], - '_impl.efc__Mgrad': out[142], - '_impl.efc__active': out[143], - '_impl.efc__alpha': out[144], - '_impl.efc__aref': out[145], - '_impl.efc__beta': out[146], - '_impl.efc__beta_den': out[147], - '_impl.efc__beta_num': out[148], - '_impl.efc__cholesky_L_tmp': out[149], - '_impl.efc__cholesky_y_tmp': out[150], - '_impl.efc__condim': out[151], - '_impl.efc__cost': out[152], - '_impl.efc__cost_candidate': out[153], - '_impl.efc__done': out[154], - '_impl.efc__force': out[155], - '_impl.efc__frictionloss': out[156], - '_impl.efc__gauss': out[157], - '_impl.efc__grad': out[158], - '_impl.efc__grad_dot': out[159], - '_impl.efc__gtol': out[160], - '_impl.efc__h': out[161], - '_impl.efc__hi': out[162], - '_impl.efc__hi_alpha': out[163], - '_impl.efc__hi_next': out[164], - '_impl.efc__hi_next_alpha': out[165], - '_impl.efc__id': out[166], - '_impl.efc__jv': out[167], - '_impl.efc__lo': out[168], - '_impl.efc__lo_alpha': out[169], - '_impl.efc__lo_next': out[170], - '_impl.efc__lo_next_alpha': out[171], - '_impl.efc__ls_done': out[172], - '_impl.efc__margin': out[173], - '_impl.efc__mid': out[174], - '_impl.efc__mid_alpha': out[175], - '_impl.efc__mv': out[176], - '_impl.efc__p0': out[177], - '_impl.efc__pos': out[178], - '_impl.efc__prev_Mgrad': out[179], - '_impl.efc__prev_cost': out[180], - '_impl.efc__prev_grad': out[181], - '_impl.efc__quad': out[182], - '_impl.efc__quad_gauss': out[183], - '_impl.efc__search': out[184], - '_impl.efc__search_dot': out[185], - '_impl.efc__type': out[186], - '_impl.efc__u': out[187], - '_impl.efc__uu': out[188], - '_impl.efc__uv': out[189], - '_impl.efc__vel': out[190], - '_impl.efc__vv': out[191], + '_impl.sensor_contact_criteria': out[94], + '_impl.sensor_contact_direction': out[95], + '_impl.sensor_contact_matchid': out[96], + '_impl.sensor_contact_nmatch': out[97], + '_impl.sensor_rangefinder_dist': out[98], + '_impl.sensor_rangefinder_geomid': out[99], + '_impl.sensor_rangefinder_pnt': out[100], + '_impl.sensor_rangefinder_vec': out[101], + 'sensordata': out[102], + 'site_xmat': out[103], + 'site_xpos': out[104], + '_impl.solver_niter': out[105], + '_impl.subtree_angmom': out[106], + '_impl.subtree_bodyvel': out[107], + 'subtree_com': out[108], + '_impl.subtree_linvel': out[109], + '_impl.ten_J': out[110], + '_impl.ten_Jdot': out[111], + '_impl.ten_actfrc': out[112], + '_impl.ten_bias_coef': out[113], + '_impl.ten_length': out[114], + '_impl.ten_velocity': out[115], + '_impl.ten_wrapadr': out[116], + '_impl.ten_wrapnum': out[117], + 'time': out[118], + '_impl.wrap_geom_xpos': out[119], + '_impl.wrap_obj': out[120], + '_impl.wrap_xpos': out[121], + 'xanchor': out[122], + 'xaxis': out[123], + 'xfrc_applied': out[124], + 'ximat': out[125], + 'xipos': out[126], + 'xmat': out[127], + 'xpos': out[128], + 'xquat': out[129], + '_impl.contact__dim': out[130], + '_impl.contact__dist': out[131], + '_impl.contact__efc_address': out[132], + '_impl.contact__frame': out[133], + '_impl.contact__friction': out[134], + '_impl.contact__geom': out[135], + '_impl.contact__includemargin': out[136], + '_impl.contact__pos': out[137], + '_impl.contact__solimp': out[138], + '_impl.contact__solref': out[139], + '_impl.contact__solreffriction': out[140], + '_impl.contact__worldid': out[141], + '_impl.efc__D': out[142], + '_impl.efc__J': out[143], + '_impl.efc__Jaref': out[144], + '_impl.efc__Ma': out[145], + '_impl.efc__Mgrad': out[146], + '_impl.efc__active': out[147], + '_impl.efc__alpha': out[148], + '_impl.efc__aref': out[149], + '_impl.efc__beta': out[150], + '_impl.efc__cholesky_L_tmp': out[151], + '_impl.efc__cholesky_y_tmp': out[152], + '_impl.efc__condim': out[153], + '_impl.efc__cost': out[154], + '_impl.efc__cost_candidate': out[155], + '_impl.efc__done': out[156], + '_impl.efc__force': out[157], + '_impl.efc__frictionloss': out[158], + '_impl.efc__gauss': out[159], + '_impl.efc__grad': out[160], + '_impl.efc__grad_dot': out[161], + '_impl.efc__gtol': out[162], + '_impl.efc__h': out[163], + '_impl.efc__hi': out[164], + '_impl.efc__hi_alpha': out[165], + '_impl.efc__hi_next': out[166], + '_impl.efc__hi_next_alpha': out[167], + '_impl.efc__id': out[168], + '_impl.efc__jv': out[169], + '_impl.efc__lo': out[170], + '_impl.efc__lo_alpha': out[171], + '_impl.efc__lo_next': out[172], + '_impl.efc__lo_next_alpha': out[173], + '_impl.efc__ls_done': out[174], + '_impl.efc__margin': out[175], + '_impl.efc__mid': out[176], + '_impl.efc__mid_alpha': out[177], + '_impl.efc__mv': out[178], + '_impl.efc__p0': out[179], + '_impl.efc__pos': out[180], + '_impl.efc__prev_Mgrad': out[181], + '_impl.efc__prev_cost': out[182], + '_impl.efc__prev_grad': out[183], + '_impl.efc__quad': out[184], + '_impl.efc__quad_gauss': out[185], + '_impl.efc__search': out[186], + '_impl.efc__search_dot': out[187], + '_impl.efc__type': out[188], + '_impl.efc__u': out[189], + '_impl.efc__uu': out[190], + '_impl.efc__uv': out[191], + '_impl.efc__vel': out[192], + '_impl.efc__vv': out[193], }) return d @@ -4114,8 +4196,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def step(m: types.Model, d: types.Data): return _step_jax_impl(m, d) - - @step.def_vmap @ffi.marshal_custom_vmap def step_vmap(unused_axis_size, is_batched, m, d): diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 087d5aab..9bb37dfb 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -42,7 +42,6 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) - @ffi.format_args_for_warp def _kinematics_shim( # Model @@ -282,8 +281,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def kinematics(m: types.Model, d: types.Data): return _kinematics_jax_impl(m, d) - - @kinematics.def_vmap @ffi.marshal_custom_vmap def kinematics_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 17cab394..52bc0f45 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -13,7 +13,6 @@ # limitations under the License. # ============================================================================== """MJX Warp types. - DO NOT EDIT. This file is auto-generated. """ import dataclasses @@ -23,10 +22,8 @@ from jax import tree_util from jax.interpreters import batching from mujoco.mjx._src import dataclasses as mjx_dataclasses import numpy as np - PyTreeNode = mjx_dataclasses.PyTreeNode - @dataclasses.dataclass(frozen=True) @tree_util.register_pytree_node_class class TileSet: @@ -38,7 +35,6 @@ class TileSet: adr: address of each tile in the set size: size of all the tiles in this set """ - adr: np.ndarray size: int @@ -59,11 +55,11 @@ class BlockDim: TODO(team): experimental and may be removed """ - actuator_velocity: int cholesky_factorize: int cholesky_factorize_solve: int cholesky_solve: int + contact_sort: int energy_vel_kinetic: int euler_dense: int mul_m_dense: int @@ -87,13 +83,10 @@ class BlockDim: class StatisticWarp(PyTreeNode): """Derived fields from Statistic.""" - meaninertia: float - class OptionWarp(PyTreeNode): """Derived fields from Option.""" - broadphase: int broadphase_filter: int epa_iterations: int @@ -106,10 +99,8 @@ class OptionWarp(PyTreeNode): sdf_initpoints: int sdf_iterations: int - class ModelWarp(PyTreeNode): """Derived fields from Model.""" - M_colind: np.ndarray M_rowadr: np.ndarray M_rownnz: np.ndarray @@ -168,6 +159,7 @@ class ModelWarp(PyTreeNode): nmeshpoly: int nmeshpolymap: int nmeshpolyvert: int + nsensortaxel: int nxn_geom_pair: np.ndarray nxn_geom_pair_filtered: np.ndarray nxn_pairid: np.ndarray @@ -183,6 +175,7 @@ class ModelWarp(PyTreeNode): qM_tiles: Tuple[TileSet, ...] rangefinder_sensor_adr: np.ndarray sensor_acc_adr: np.ndarray + sensor_contact_adr: np.ndarray sensor_e_kinetic: bool sensor_e_potential: bool sensor_limitfrc_adr: np.ndarray @@ -197,6 +190,8 @@ class ModelWarp(PyTreeNode): sensor_touch_adr: np.ndarray sensor_vel_adr: np.ndarray subtree_mass: jax.Array + taxel_sensorid: np.ndarray + taxel_vertadr: np.ndarray ten_wrapadr_site: np.ndarray ten_wrapnum_site: np.ndarray tendon_geom_adr: np.ndarray @@ -209,10 +204,8 @@ class ModelWarp(PyTreeNode): wrap_site_adr: np.ndarray wrap_site_pair_adr: np.ndarray - class DataWarp(PyTreeNode): """Derived fields from Data.""" - act_dot_rk: jax.Array act_t0: jax.Array act_vel_integration: jax.Array @@ -252,8 +245,6 @@ class DataWarp(PyTreeNode): efc__alpha: jax.Array efc__aref: jax.Array efc__beta: jax.Array - efc__beta_den: jax.Array - efc__beta_num: jax.Array efc__cholesky_L_tmp: jax.Array efc__cholesky_y_tmp: jax.Array efc__condim: jax.Array @@ -321,6 +312,7 @@ class DataWarp(PyTreeNode): ncollision: jax.Array ncon: jax.Array ncon_hfield: jax.Array + ncon_world: jax.Array nconmax: int ne: jax.Array ne_connect: jax.Array @@ -357,6 +349,10 @@ class DataWarp(PyTreeNode): sap_range: jax.Array sap_segment_index: jax.Array sap_sort_index: jax.Array + sensor_contact_criteria: jax.Array + sensor_contact_direction: jax.Array + sensor_contact_matchid: jax.Array + sensor_contact_nmatch: jax.Array sensor_rangefinder_dist: jax.Array sensor_rangefinder_geomid: jax.Array sensor_rangefinder_pnt: jax.Array @@ -377,8 +373,6 @@ class DataWarp(PyTreeNode): wrap_obj: jax.Array wrap_xpos: jax.Array shape = property(lambda self: self.cacc.shape) - - DATA_NON_VMAP = { 'collision_hftri_index', 'collision_pair', @@ -421,7 +415,6 @@ DATA_NON_VMAP = { 'ray_bodyexclude', } - def _to_elt(cont, _, d, axis): return DataWarp(**{ f.name: ( @@ -494,8 +487,6 @@ NDIM = { 'efc__alpha': 1, 'efc__aref': 2, 'efc__beta': 1, - 'efc__beta_den': 1, - 'efc__beta_num': 1, 'efc__cholesky_L_tmp': 3, 'efc__cholesky_y_tmp': 2, 'efc__condim': 2, @@ -568,6 +559,7 @@ NDIM = { 'ncollision': 1, 'ncon': 1, 'ncon_hfield': 2, + 'ncon_world': 1, 'nconmax': 0, 'ne': 1, 'ne_connect': 1, @@ -618,6 +610,10 @@ NDIM = { 'sap_range': 2, 'sap_segment_index': 2, 'sap_sort_index': 3, + 'sensor_contact_criteria': 3, + 'sensor_contact_direction': 3, + 'sensor_contact_matchid': 3, + 'sensor_contact_nmatch': 2, 'sensor_rangefinder_dist': 2, 'sensor_rangefinder_geomid': 2, 'sensor_rangefinder_pnt': 3, @@ -684,6 +680,7 @@ NDIM = { 'block_dim__cholesky_factorize': 0, 'block_dim__cholesky_factorize_solve': 0, 'block_dim__cholesky_solve': 0, + 'block_dim__contact_sort': 0, 'block_dim__energy_vel_kinetic': 0, 'block_dim__euler_dense': 0, 'block_dim__mul_m_dense': 0, @@ -829,6 +826,8 @@ NDIM = { 'mesh_faceadr': 1, 'mesh_graph': 1, 'mesh_graphadr': 1, + 'mesh_normal': 2, + 'mesh_normaladr': 1, 'mesh_polyadr': 1, 'mesh_polymap': 1, 'mesh_polymapadr': 1, @@ -838,6 +837,7 @@ NDIM = { 'mesh_polyvert': 1, 'mesh_polyvertadr': 1, 'mesh_polyvertnum': 1, + 'mesh_quat': 2, 'mesh_vert': 2, 'mesh_vertadr': 1, 'mesh_vertnum': 1, @@ -872,6 +872,7 @@ NDIM = { 'nq': 0, 'nsensor': 0, 'nsensordata': 0, + 'nsensortaxel': 0, 'nsite': 0, 'ntendon': 0, 'nu': 0, @@ -931,11 +932,13 @@ NDIM = { 'rangefinder_sensor_adr': 1, 'sensor_acc_adr': 1, 'sensor_adr': 1, + 'sensor_contact_adr': 1, 'sensor_cutoff': 1, 'sensor_datatype': 1, 'sensor_dim': 1, 'sensor_e_kinetic': 0, 'sensor_e_potential': 0, + 'sensor_intprm': 2, 'sensor_limitfrc_adr': 1, 'sensor_limitpos_adr': 1, 'sensor_limitvel_adr': 1, @@ -959,6 +962,8 @@ NDIM = { 'site_type': 1, 'stat__meaninertia': 0, 'subtree_mass': 2, + 'taxel_sensorid': 1, + 'taxel_vertadr': 1, 'ten_wrapadr_site': 1, 'ten_wrapnum_site': 1, 'tendon_actfrclimited': 1, @@ -1071,8 +1076,6 @@ BATCH_DIM = { 'efc__alpha': True, 'efc__aref': True, 'efc__beta': True, - 'efc__beta_den': True, - 'efc__beta_num': True, 'efc__cholesky_L_tmp': True, 'efc__cholesky_y_tmp': True, 'efc__condim': True, @@ -1145,6 +1148,7 @@ BATCH_DIM = { 'ncollision': False, 'ncon': False, 'ncon_hfield': True, + 'ncon_world': True, 'nconmax': False, 'ne': True, 'ne_connect': True, @@ -1195,6 +1199,10 @@ BATCH_DIM = { 'sap_range': True, 'sap_segment_index': True, 'sap_sort_index': True, + 'sensor_contact_criteria': True, + 'sensor_contact_direction': True, + 'sensor_contact_matchid': True, + 'sensor_contact_nmatch': True, 'sensor_rangefinder_dist': True, 'sensor_rangefinder_geomid': True, 'sensor_rangefinder_pnt': True, @@ -1261,6 +1269,7 @@ BATCH_DIM = { 'block_dim__cholesky_factorize': False, 'block_dim__cholesky_factorize_solve': False, 'block_dim__cholesky_solve': False, + 'block_dim__contact_sort': False, 'block_dim__energy_vel_kinetic': False, 'block_dim__euler_dense': False, 'block_dim__mul_m_dense': False, @@ -1406,6 +1415,8 @@ BATCH_DIM = { 'mesh_faceadr': False, 'mesh_graph': False, 'mesh_graphadr': False, + 'mesh_normal': False, + 'mesh_normaladr': False, 'mesh_polyadr': False, 'mesh_polymap': False, 'mesh_polymapadr': False, @@ -1415,6 +1426,7 @@ BATCH_DIM = { 'mesh_polyvert': False, 'mesh_polyvertadr': False, 'mesh_polyvertnum': False, + 'mesh_quat': False, 'mesh_vert': False, 'mesh_vertadr': False, 'mesh_vertnum': False, @@ -1449,6 +1461,7 @@ BATCH_DIM = { 'nq': False, 'nsensor': False, 'nsensordata': False, + 'nsensortaxel': False, 'nsite': False, 'ntendon': False, 'nu': False, @@ -1508,11 +1521,13 @@ BATCH_DIM = { 'rangefinder_sensor_adr': False, 'sensor_acc_adr': False, 'sensor_adr': False, + 'sensor_contact_adr': False, 'sensor_cutoff': False, 'sensor_datatype': False, 'sensor_dim': False, 'sensor_e_kinetic': False, 'sensor_e_potential': False, + 'sensor_intprm': False, 'sensor_limitfrc_adr': False, 'sensor_limitpos_adr': False, 'sensor_limitvel_adr': False, @@ -1536,6 +1551,8 @@ BATCH_DIM = { 'site_type': False, 'stat__meaninertia': False, 'subtree_mass': True, + 'taxel_sensorid': False, + 'taxel_vertadr': False, 'ten_wrapadr_site': False, 'ten_wrapnum_site': False, 'tendon_actfrclimited': False,