From a359780ca16e5c8109fa877054709aff9b845c86 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Fri, 21 Nov 2025 08:15:42 -0800 Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub. PiperOrigin-RevId: 835230315 Change-Id: Ifa4f8f66c4e3de1a3074f0897daf93f88e469f06 --- .../mujoco_warp/_src/collision_convex.py | 346 +++- .../mujoco_warp/_src/collision_driver.py | 42 +- .../mujoco_warp/_src/collision_gjk.py | 54 +- .../mujoco_warp/_src/collision_primitive.py | 15 +- .../mujoco_warp/_src/collision_sdf.py | 16 +- .../mujoco_warp/_src/constraint.py | 4 +- .../mujoco_warp/_src/derivative.py | 21 +- .../third_party/mujoco_warp/_src/forward.py | 51 +- .../mjx/third_party/mujoco_warp/_src/io.py | 1750 +++++------------ .../third_party/mujoco_warp/_src/passive.py | 14 +- .../third_party/mujoco_warp/_src/sensor.py | 10 +- .../third_party/mujoco_warp/_src/smooth.py | 2 +- .../third_party/mujoco_warp/_src/solver.py | 13 +- .../third_party/mujoco_warp/_src/support.py | 2 +- .../mjx/third_party/mujoco_warp/_src/types.py | 1151 +++++------ .../third_party/mujoco_warp/_src/warp_util.py | 12 +- mjx/mujoco/mjx/warp/collision_driver.py | 2 +- mjx/mujoco/mjx/warp/forward.py | 30 +- mjx/mujoco/mjx/warp/forward_test.py | 3 +- mjx/mujoco/mjx/warp/types.py | 44 +- 20 files changed, 1578 insertions(+), 2004 deletions(-) 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 99de4580..d3772f89 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 @@ -32,17 +32,18 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import Model +from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 +from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) -MULTI_CONTACT_COUNT = 8 -mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float) -mat63 = wp.types.matrix(shape=(6, 3), dtype=float) +vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float) +mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float) _CONVEX_COLLISION_PAIRS = [ (GeomType.HFIELD, GeomType.SPHERE), @@ -260,9 +261,9 @@ def ccd_kernel_builder( contact_geomcollisionid_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int), ) -> int: - points = mat3c() - witness1 = mat3c() - witness2 = mat3c() + points = mat43() + witness1 = mat43() + witness2 = mat43() geom1.margin = margin geom2.margin = margin if pairid[1] >= 0: @@ -566,6 +567,33 @@ def ccd_kernel_builder( adr = hfield_adr[geom1_dataid] + hfield_contact_dist = vec_maxconpair() + hfield_contact_pos = mat_maxconpair() + hfield_contact_normal = mat_maxconpair() + min_dist = float(wp.inf) + min_normal = wp.vec3(wp.inf, wp.inf, wp.inf) + min_pos = wp.vec3(wp.inf, wp.inf, wp.inf) + min_id = int(-1) + + # TODO(team): height field margin? + geom1.margin = margin + geom2.margin = margin + + # EPA memory + epa_vert = epa_vert_in[tid] + epa_vert1 = epa_vert1_in[tid] + epa_vert2 = epa_vert2_in[tid] + epa_vert_index1 = epa_vert_index1_in[tid] + epa_vert_index2 = epa_vert_index2_in[tid] + epa_face = epa_face_in[tid] + epa_pr = epa_pr_in[tid] + epa_norm2 = epa_norm2_in[tid] + epa_index = epa_index_in[tid] + epa_map = epa_map_in[tid] + epa_horizon = epa_horizon_in[tid] + + collision_pairid = collision_pairid_in[tid] + # process all prisms in subgrid count = int(0) for r in range(rmin, rmax): @@ -573,6 +601,13 @@ def ccd_kernel_builder( for c in range(cmin, cmax + 1): # add both triangles from this cell for i in range(2): + if count >= MJ_MAXCONPAIR: + wp.printf( + "height field collision overflow, number of collisions >= %u - please adjust resolution: \n decrease the number of hfield rows/cols or modify size of colliding geom\n", + MJ_MAXCONPAIR, + ) + continue + # add vert x = dx * float(c) - size[0] y = dy * float(r + dr[i]) - size[1] @@ -607,66 +642,253 @@ def ccd_kernel_builder( x1_ += prism[i] x1 += geom1.rot @ (x1_ / 6.0) - ncontact = eval_ccd_write_contact( - opt_ccd_tolerance, - geom_type, - naconmax_in, - epa_vert_in, - epa_vert1_in, - epa_vert2_in, - epa_vert_index1_in, - epa_vert_index2_in, - epa_face_in, - epa_pr_in, - epa_norm2_in, - epa_index_in, - epa_map_in, - epa_horizon_in, - multiccd_polygon_in, - multiccd_clipped_in, - multiccd_pnormal_in, - multiccd_pdist_in, - multiccd_idx1_in, - multiccd_idx2_in, - multiccd_n1_in, - multiccd_n2_in, - multiccd_endvert_in, - multiccd_face1_in, - multiccd_face2_in, + dist, ncontact, w1, w2, idx = ccd( + opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], + 0.0, + ccd_iterations, geom1, geom2, - geoms, - worldid, - tid, - margin, - gap, - condim, - friction, - solref, - solreffriction, - solimp, + geomtype1, + geomtype2, x1, geom2.pos, - count, - collision_pairid_in[tid], - contact_dist_out, - contact_pos_out, - contact_frame_out, - contact_includemargin_out, - contact_friction_out, - contact_solref_out, - contact_solreffriction_out, - contact_solimp_out, - contact_dim_out, - contact_geom_out, - contact_worldid_out, - contact_type_out, - contact_geomcollisionid_out, - nacon_out, + epa_vert, + epa_vert1, + epa_vert2, + epa_vert_index1, + epa_vert_index2, + epa_face, + epa_pr, + epa_norm2, + epa_index, + epa_map, + epa_horizon, ) - count += ncontact - if count >= MJ_MAXCONPAIR: - return + + if ncontact == 0: + continue + + # cache contact information + hfield_contact_dist[count] = dist + + pos = 0.5 * (w1 + w2) + hfield_contact_pos[count, 0] = pos[0] + hfield_contact_pos[count, 1] = pos[1] + hfield_contact_pos[count, 2] = pos[2] + + frame = make_frame(w1 - w2) + normal = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2]) + hfield_contact_normal[count, 0] = normal[0] + hfield_contact_normal[count, 1] = normal[1] + hfield_contact_normal[count, 2] = normal[2] + + # contact with minimum distance + if dist < min_dist: + min_dist = dist + min_normal = normal + min_pos = pos + min_id = count + + count += 1 + + # contact 0: minimum distance + write_contact( + naconmax_in, + 0, + min_dist, + min_pos, + make_frame(min_normal), + margin, + gap, + condim, + friction, + solref, + solreffriction, + solimp, + geoms, + collision_pairid, + worldid, + contact_dist_out, + contact_pos_out, + contact_frame_out, + contact_includemargin_out, + contact_friction_out, + contact_solref_out, + contact_solreffriction_out, + contact_solimp_out, + contact_dim_out, + contact_geom_out, + contact_worldid_out, + contact_type_out, + contact_geomcollisionid_out, + nacon_out, + ) + + # TODO(team): routine for select subset of contacts + # TODO(team): if use_multiccd? + if wp.static(True): + MIN_DIST_TO_NEXT_CONTACT = 1.0e-3 + + # contact 1: furthest from minimum distance contact + id1 = int(-1) + dist1 = float(-wp.inf) + for i in range(count): + if i == min_id: + continue + + hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2]) + dist = wp.norm_l2(hf_pos - min_pos) + + if dist > dist1: + id1 = i + dist1 = dist + + if id1 == -1 or (0.0 < dist1 and dist1 < MIN_DIST_TO_NEXT_CONTACT): + return + + pos1 = wp.vec3(hfield_contact_pos[id1, 0], hfield_contact_pos[id1, 1], hfield_contact_pos[id1, 2]) + normal1 = wp.vec3(hfield_contact_normal[id1, 0], hfield_contact_normal[id1, 1], hfield_contact_normal[id1, 2]) + + write_contact( + naconmax_in, + 1, + hfield_contact_dist[id1], + pos1, + make_frame(normal1), + margin, + gap, + condim, + friction, + solref, + solreffriction, + solimp, + geoms, + collision_pairid, + worldid, + contact_dist_out, + contact_pos_out, + contact_frame_out, + contact_includemargin_out, + contact_friction_out, + contact_solref_out, + contact_solreffriction_out, + contact_solimp_out, + contact_dim_out, + contact_geom_out, + contact_worldid_out, + contact_type_out, + contact_geomcollisionid_out, + nacon_out, + ) + + # contact 2: point furthest from min_pos - pos1 line + dist_min1 = wp.cross(min_normal, min_pos - pos1) + + id2 = int(-1) + dist_12 = float(-wp.inf) + for i in range(count): + if i == min_id or i == id1: + continue + + hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2]) + dist = wp.abs(wp.dot(hf_pos - min_pos, dist_min1)) + + if dist > dist_12: + id2 = i + dist_12 = dist + + if id2 == -1 or (0.0 < dist_12 and dist_12 < MIN_DIST_TO_NEXT_CONTACT): + return + + pos2 = wp.vec3(hfield_contact_pos[id2, 0], hfield_contact_pos[id2, 1], hfield_contact_pos[id2, 2]) + normal2 = wp.vec3(hfield_contact_normal[id2, 0], hfield_contact_normal[id2, 1], hfield_contact_normal[id2, 2]) + + write_contact( + naconmax_in, + 2, + hfield_contact_dist[id2], + pos2, + make_frame(normal2), + margin, + gap, + condim, + friction, + solref, + solreffriction, + solimp, + geoms, + collision_pairid, + worldid, + contact_dist_out, + contact_pos_out, + contact_frame_out, + contact_includemargin_out, + contact_friction_out, + contact_solref_out, + contact_solreffriction_out, + contact_solimp_out, + contact_dim_out, + contact_geom_out, + contact_worldid_out, + contact_type_out, + contact_geomcollisionid_out, + nacon_out, + ) + + # contact 3: point furthest from other triangle edge + vec_min2 = wp.cross(min_normal, min_pos - pos2) + vec_12 = wp.cross(min_normal, pos1 - pos2) + + id3 = int(-1) + dist3 = float(-wp.inf) + for i in range(count): + if i == min_id or i == id1 or i == id2: + continue + + hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2]) + dist = wp.abs(wp.dot(hf_pos - min_pos, vec_min2)) + wp.abs(wp.dot(pos1 - hf_pos, vec_12)) + + if dist > dist3: + id3 = i + dist3 = dist + + if id3 == -1 or (0.0 < dist3 and dist3 < MIN_DIST_TO_NEXT_CONTACT): + return + + pos3 = wp.vec3(hfield_contact_pos[id3, 0], hfield_contact_pos[id3, 1], hfield_contact_pos[id3, 2]) + normal3 = wp.vec3(hfield_contact_normal[id3, 0], hfield_contact_normal[id3, 1], hfield_contact_normal[id3, 2]) + + write_contact( + naconmax_in, + 3, + hfield_contact_dist[id3], + pos3, + make_frame(normal3), + margin, + gap, + condim, + friction, + solref, + solreffriction, + solimp, + geoms, + collision_pairid, + worldid, + contact_dist_out, + contact_pos_out, + contact_frame_out, + contact_includemargin_out, + contact_friction_out, + contact_solref_out, + contact_solreffriction_out, + contact_solimp_out, + contact_dim_out, + contact_geom_out, + contact_worldid_out, + contact_type_out, + contact_geomcollisionid_out, + nacon_out, + ) else: eval_ccd_write_contact( opt_ccd_tolerance, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py index ef314275..60df7d3a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py @@ -27,9 +27,11 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import Model +from mujoco.mjx.third_party.mujoco_warp._src.types import mat23 +from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) @@ -170,10 +172,6 @@ def _aabb_filter( return True -mat23 = wp.types.matrix(shape=(2, 3), dtype=float) -mat63 = wp.types.matrix(shape=(6, 3), dtype=float) - - # TODO(team): improve performance by precomputing bounding box @wp.func def _obb_filter( @@ -232,7 +230,7 @@ def _obb_filter( return True -def _broadphase_filter(m: Model): +def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int): @wp.func def func( # Model: @@ -252,28 +250,28 @@ def _broadphase_filter(m: Model): # 4: aabb # 8: obb - aabb_id = worldid % geom_aabb.shape[0] if wp.static(m.geom_aabb.shape[0] > 1) else 0 + aabb_id = worldid % ngeom_aabb if wp.static(ngeom_aabb > 1) else 0 center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0] size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1] - rbound_id = worldid % geom_rbound.shape[0] if wp.static(m.geom_rbound.shape[0] > 1) else 0 + rbound_id = worldid % ngeom_rbound if wp.static(ngeom_rbound > 1) else 0 rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2] - margin_id = worldid % geom_margin.shape[0] if wp.static(m.geom_margin.shape[0] > 1) else 0 + margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0 margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2] xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2] if rbound1 == 0.0 or rbound2 == 0.0: - if wp.static(m.opt.broadphase_filter & BroadphaseFilter.PLANE): + if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE): return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2) else: - if wp.static(m.opt.broadphase_filter & BroadphaseFilter.SPHERE): + if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE): if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2): return False - if wp.static(m.opt.broadphase_filter & BroadphaseFilter.AABB): + if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB): if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): return False - if wp.static(m.opt.broadphase_filter & BroadphaseFilter.OBB): + if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB): if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): return False @@ -403,7 +401,7 @@ def _sap_range( @cache_kernel -def _sap_broadphase(broadphase_filter): +def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int): @nested_kernel(module="unique", enable_backward=False) def kernel( # Model: @@ -461,7 +459,9 @@ def _sap_broadphase(broadphase_filter): continue if ( - broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid) + wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))( + geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid + ) or pairid[1] >= 0 ): _add_geom_pair( @@ -578,9 +578,8 @@ def sap_broadphase(m: Model, d: Data): # estimate number of overlap checks # assumes each geom has 5 other geoms (batched over all worlds) nsweep = 5 * nworldgeom - broadphase_filter = _broadphase_filter(m) wp.launch( - kernel=_sap_broadphase(broadphase_filter), + kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]), dim=nsweep, inputs=[ m.ngeom, @@ -602,7 +601,7 @@ def sap_broadphase(m: Model, d: Data): @cache_kernel -def _nxn_broadphase(broadphase_filter): +def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int): @nested_kernel(module="unique", enable_backward=False) def kernel( # Model: @@ -629,7 +628,9 @@ def _nxn_broadphase(broadphase_filter): geom2 = geom[1] if ( - broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid) + wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))( + geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid + ) or nxn_pairid[elementid][1] >= 0 ): _add_geom_pair( @@ -663,9 +664,8 @@ def nxn_broadphase(m: Model, d: Data): The initial list of pairs is filtered at model creation time to exclude pairs based on `contype`/`conaffinity`, parent-child relationships, and explicit `` tags. """ - broadphase_filter = _broadphase_filter(m) wp.launch( - _nxn_broadphase(broadphase_filter), + _nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]), dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]), inputs=[ m.geom_type, 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 bd0b29bb..8d59525a 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 @@ -20,6 +20,8 @@ import warp as wp from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType +from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 +from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) @@ -32,13 +34,6 @@ MJ_MINVAL2 = MJ_MINVAL * MJ_MINVAL FACE_TOL = 0.99999872 EDGE_TOL = 0.00159999931 -mat43 = wp.types.matrix(shape=(4, 3), dtype=float) -mat63 = wp.types.matrix(shape=(6, 3), dtype=float) - -# box-box currently supports up to 8 contacts -MULTI_CONTACT_COUNT = 8 -mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float) - @wp.struct class GJKResult: @@ -155,30 +150,25 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: vert_edgeadr = geom.graphadr + 2 vert_globalid = geom.graphadr + 2 + numvert edge_localid = geom.graphadr + 2 + 2 * numvert - # hillclimb until no change prev = int(-1) - imax = int(0) - if geom.index > -1: - imax = geom.index - sp.cached_index = geom.index + imax = wp.where(geom.index > -1, geom.index, 0) - while True: - prev = int(imax) - i = int(geom.graph[vert_edgeadr + imax]) - while geom.graph[edge_localid + i] >= 0: - subidx = geom.graph[edge_localid + i] + # hillclimb until no change + while imax != prev: + prev = imax + i = geom.graph[vert_edgeadr + imax] + subidx = geom.graph[edge_localid + i] + while subidx >= 0: idx = geom.graph[vert_globalid + subidx] dist = wp.dot(local_dir, geom.vert[geom.vertadr + idx]) - if dist > max_dist: - max_dist = dist - imax = int(subidx) - i += int(1) - if imax == prev: - break + imax = wp.where(dist > max_dist, subidx, imax) + max_dist = wp.where(dist > max_dist, dist, max_dist) + i += 1 + subidx = geom.graph[edge_localid + i] + sp.cached_index = imax - imax = geom.graph[vert_globalid + imax] - sp.vertex_index = imax - sp.point = geom.vert[geom.vertadr + imax] + sp.vertex_index = geom.graph[vert_globalid + imax] + sp.point = geom.vert[geom.vertadr + sp.vertex_index] sp.point = geom.rot @ sp.point + geom.pos elif geomtype == GeomType.HFIELD: @@ -1813,9 +1803,9 @@ def _polygon_clip( # Out: polygon_out: wp.array(dtype=wp.vec3), clipped_out: wp.array(dtype=wp.vec3), -) -> Tuple[int, mat3c, mat3c]: - witness1 = mat3c() - witness2 = mat3c() +) -> Tuple[int, mat43, mat43]: + witness1 = mat43() + witness2 = mat43() # clipping face needs to be at least a triangle if nface1 < 3: @@ -1930,9 +1920,9 @@ def multicontact( geom2: Geom, geomtype1: int, geomtype2: int, -) -> Tuple[int, mat3c, mat3c]: - witness1 = mat3c() - witness2 = mat3c() +) -> Tuple[int, mat43, mat43]: + witness1 = mat43() + witness2 = mat43() witness1[0] = x1 witness2[0] = x2 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py index bd17475f..4412c4cc 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py @@ -38,21 +38,16 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import Model +from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 +from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) -class mat43f(wp.types.matrix(shape=(4, 3), dtype=wp.float32)): - pass - - -mat63 = wp.types.matrix(shape=(6, 3), dtype=float) - - @wp.struct class Geom: pos: wp.vec3 @@ -169,7 +164,7 @@ def geom_collision_pair( @wp.func -def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43f, wp.vec3]: +def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43, wp.vec3]: """Core contact geometry calculation for plane-convex collision. Args: @@ -185,7 +180,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup _HUGE_VAL = 1e6 contact_dist = wp.vec4(wp.inf) - contact_pos = mat43f() + contact_pos = mat43() contact_count = int(0) # get points in the convex frame 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 93b607f9..4e855036 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 @@ -26,7 +26,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 -from mujoco.mjx.third_party.mujoco_warp._src.types import vec8f +from mujoco.mjx.third_party.mujoco_warp._src.types import vec8 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope @@ -54,7 +54,7 @@ class VolumeData: half_size: wp.vec3 oct_aabb: wp.array2d(dtype=wp.vec3) oct_child: wp.array(dtype=vec8i) - oct_coeff: wp.array(dtype=vec8f) + oct_coeff: wp.array(dtype=vec8) valid: bool = False @@ -79,7 +79,7 @@ def get_sdf_params( # Model: oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8f), + oct_coeff: wp.array(dtype=vec8), plugin: wp.array(dtype=int), plugin_attr: wp.array(dtype=wp.vec3f), # In: @@ -225,12 +225,12 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3: @wp.func def find_oct( oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool -) -> Tuple[int, Tuple[vec8f, vec8f, vec8f]]: +) -> Tuple[int, Tuple[vec8, vec8, vec8]]: stack = int(0) niter = int(100) - rx = vec8f(0.0) - ry = vec8f(0.0) - rz = vec8f(0.0) + rx = vec8(0.0) + ry = vec8(0.0) + rz = vec8(0.0) eps = 1e-6 while niter > 0: @@ -621,7 +621,7 @@ def _sdf_narrowphase( nmeshface: int, oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8f), + oct_coeff: wp.array(dtype=vec8), geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py index f0fc7a30..831565f9 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py @@ -1843,7 +1843,7 @@ def make_constraint(m: types.Model, d: types.Data): if m.opt.cone == types.ConeType.PYRAMIDAL: wp.launch( _efc_contact_pyramidal, - dim=(d.naconmax, 2 * (m.condim_max - 1) if m.condim_max > 1 else 1), + dim=(d.naconmax, m.nmaxpyramid), inputs=[ m.nv, m.opt.timestep, @@ -1888,7 +1888,7 @@ def make_constraint(m: types.Model, d: types.Data): elif m.opt.cone == types.ConeType.ELLIPTIC: wp.launch( _efc_contact_elliptic, - dim=(d.naconmax, m.condim_max), + dim=(d.naconmax, m.nmaxcondim), inputs=[ m.nv, m.opt.timestep, diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py index 13a5dbf2..83332751 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py @@ -75,14 +75,21 @@ def _qderiv_actuator_passive( else: bias = 0.0 - if actuator_dyntype[actid] != DynType.NONE: - act_first = actuator_actadr[actid] - act_last = act_first + actuator_actnum[actid] - 1 - vel = bias + gain * act_in[worldid, act_last] - else: - vel = bias + gain * ctrl_in[worldid, actid] + if bias == 0.0 and gain == 0.0: + continue - qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel + vel = bias + if actuator_dyntype[actid] != DynType.NONE: + if gain != 0.0: + act_first = actuator_actadr[actid] + act_last = act_first + actuator_actnum[actid] - 1 + vel += gain * act_in[worldid, act_last] + else: + if gain != 0.0: + vel += gain * ctrl_in[worldid, actid] + + if vel != 0.0: + qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel # TODO(team): fluid model derivative 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 87f50557..41df144d 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -42,7 +42,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) @@ -521,7 +521,8 @@ def fwd_position(m: Model, d: Data, factorize: bool = True): # TODO(team): sparse actuator_moment version -def _actuator_velocity(m: Model, d: Data): +@cache_kernel +def _actuator_velocity(nv: int): @nested_kernel(module="unique", enable_backward=False) def actuator_velocity( # Data in: @@ -531,22 +532,17 @@ def _actuator_velocity(m: Model, d: Data): actuator_velocity_out: wp.array2d(dtype=float), ): worldid, actid = wp.tid() - moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(m.nv)) - qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(m.nv)) + moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(nv)) + qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv)) moment_qvel_tile = wp.tile_map(wp.mul, moment_tile, qvel_tile) actuator_velocity_tile = wp.tile_reduce(wp.add, moment_qvel_tile) actuator_velocity_out[worldid, actid] = actuator_velocity_tile[0] - wp.launch_tiled( - actuator_velocity, - dim=(d.nworld, m.nu), - inputs=[d.qvel, d.actuator_moment], - outputs=[d.actuator_velocity], - block_dim=m.block_dim.actuator_velocity, - ) + return actuator_velocity -def _tendon_velocity(m: Model, d: Data): +@cache_kernel +def _tendon_velocity(nv: int): @nested_kernel(module="unique", enable_backward=False) def tendon_velocity( # Data in: @@ -556,29 +552,34 @@ def _tendon_velocity(m: Model, d: Data): ten_velocity_out: wp.array2d(dtype=float), ): worldid, tenid = wp.tid() - ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(m.nv)) - qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(m.nv)) + ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(nv)) + qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv)) ten_J_qvel_tile = wp.tile_map(wp.mul, ten_J_tile, qvel_tile) ten_velocity_tile = wp.tile_reduce(wp.add, ten_J_qvel_tile) ten_velocity_out[worldid, tenid] = ten_velocity_tile[0] - wp.launch_tiled( - tendon_velocity, - dim=(d.nworld, m.ntendon), - inputs=[d.qvel, d.ten_J], - outputs=[d.ten_velocity], - block_dim=m.block_dim.tendon_velocity, - ) + return tendon_velocity @event_scope def fwd_velocity(m: Model, d: Data): """Velocity-dependent computations.""" - _actuator_velocity(m, d) + wp.launch_tiled( + _actuator_velocity(m.nv), + dim=(d.nworld, m.nu), + inputs=[d.qvel, d.actuator_moment], + outputs=[d.actuator_velocity], + block_dim=m.block_dim.actuator_velocity, + ) - if m.ntendon > 0: - # TODO(team): sparse version - _tendon_velocity(m, d) + # TODO(team): sparse version + wp.launch_tiled( + _tendon_velocity(m.nv), + dim=(d.nworld, m.ntendon), + inputs=[d.qvel, d.ten_J], + outputs=[d.ten_velocity], + block_dim=m.block_dim.tendon_velocity, + ) smooth.com_vel(m, d) passive.passive(m, d) 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 4e14734f..2af891c5 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -13,59 +13,40 @@ # limitations under the License. # ============================================================================== +import dataclasses from typing import Any, Optional, Sequence, Union import mujoco import numpy as np import warp as wp -from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import types from mujoco.mjx.third_party.mujoco_warp._src import warp_util -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel - -# tolerance override for float32 -_TOLERANCE_F32 = 1.0e-6 +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel -def _compute_nmaxpolygon(mjm: mujoco.MjModel, geom_pair_type_count: tuple[int, ...]) -> int: - """Compute nmaxpolygon given the geom pairs for the model.""" - nboxbox = geom_pair_type_count[math.upper_trid_index(len(types.GeomType), types.GeomType.BOX.value, types.GeomType.BOX.value)] - nboxmesh = geom_pair_type_count[ - math.upper_trid_index(len(types.GeomType), types.GeomType.BOX.value, types.GeomType.MESH.value) - ] - nmeshmesh = geom_pair_type_count[ - math.upper_trid_index(len(types.GeomType), types.GeomType.MESH.value, types.GeomType.MESH.value) - ] +def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> Union[wp.array, None]: + """Creates a warp array and populates it with data. - # need at least 4 (square sides) if there's a box collision needing multiccd - # TODO(kbayes): remove nboxbox or enable ccd for box-box collisions - box_factor = 4 if nboxbox + nboxmesh > 0 else 0 + The array shape is determined by a field spec referencing MjModel / MjData array sizes. + """ + shape = None + if spec.shape != (0,): + shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec.shape) - # possibly need to allocate more memory if there's meshes - if nmeshmesh + nboxmesh > 0: - return np.append(mjm.mesh_polyvertnum, box_factor).max() - return box_factor + if data is None and shape is None: + return None # nothing to do + elif data is None: + array = wp.zeros(shape, dtype=spec.dtype) + else: + array = wp.array(np.array(data), dtype=spec.dtype, shape=shape) - -def _compute_nmaxmeshdeg(mjm: mujoco.MjModel, geom_pair_type_count: tuple[int, ...]) -> int: - """Compute nmaxmeshdeg given the geom pairs for the model.""" - nboxbox = geom_pair_type_count[math.upper_trid_index(len(types.GeomType), types.GeomType.BOX.value, types.GeomType.BOX.value)] - nboxmesh = geom_pair_type_count[ - math.upper_trid_index(len(types.GeomType), types.GeomType.BOX.value, types.GeomType.MESH.value) - ] - nmeshmesh = geom_pair_type_count[ - math.upper_trid_index(len(types.GeomType), types.GeomType.MESH.value, types.GeomType.MESH.value) - ] - - # need at least 3 (3 edges per vertex) if there's a box collision needing multiccd - # TODO(kbayes): remove nboxbox or enable ccd for box-box collisions - box_factor = 3 if nboxbox + nboxmesh > 0 else 0 - - # possibly need to allocate more memory if there's meshes - if nmeshmesh + nboxmesh > 0: - return np.append(mjm.mesh_polymapnum, box_factor).max() - return box_factor + if spec.shape[0] == "*": + # add private attribute for JAX to determine which fields are batched + array._is_batched = True + # also set stride 0 to 0 which is expected legacy behavior (but is deprecated) + array.strides = (0,) + array.strides[1:] + return array def put_model(mjm: mujoco.MjModel) -> types.Model: @@ -80,56 +61,37 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: # check for compatible cuda toolkit and driver versions warp_util.check_toolkit_driver() - # check supported features - for field, field_types, field_str in ( - (mjm.actuator_trntype, types.TrnType, "Actuator transmission type"), - (mjm.actuator_dyntype, types.DynType, "Actuator dynamics type"), - (mjm.actuator_gaintype, types.GainType, "Gain type"), - (mjm.actuator_biastype, types.BiasType, "Bias type"), - (mjm.eq_type, types.EqType, "Equality constraint types"), - (mjm.geom_type, types.GeomType, "Geom type"), - (mjm.sensor_type, types.SensorType, "Sensor types"), - (mjm.wrap_type, types.WrapType, "Wrap types"), + # model: check supported features in array types + for field, field_type in ( + (mjm.actuator_trntype, types.TrnType), + (mjm.actuator_dyntype, types.DynType), + (mjm.actuator_gaintype, types.GainType), + (mjm.actuator_biastype, types.BiasType), + (mjm.eq_type, types.EqType), + (mjm.geom_type, types.GeomType), + (mjm.sensor_type, types.SensorType), + (mjm.wrap_type, types.WrapType), ): - unsupported = ~np.isin(field, list(field_types)) - if unsupported.any(): - raise NotImplementedError(f"{field_str} {field[unsupported]} not supported.") + missing = ~np.isin(field, field_type) + if missing.any(): + raise NotImplementedError(f"{field_type.__name__}: {field[missing]} not supported.") - plugin_id = [] - plugin_attr = [] - geom_plugin_index = np.full_like(mjm.geom_type, -1) + # opt: check supported features in scalar types + for field, field_type in ( + (mjm.opt.integrator, types.IntegratorType), + (mjm.opt.cone, types.ConeType), + (mjm.opt.solver, types.SolverType), + ): + if field not in set(field_type): + raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.") - if mjm.nplugin > 0: - if (mjm.body_plugin != -1).any(): - raise NotImplementedError("Body plugins not supported.") - if (mjm.actuator_plugin != -1).any(): - raise NotImplementedError("Actuator plugins not supported.") - if (mjm.sensor_plugin != -1).any(): - raise NotImplementedError("Sensor plugins not supported.") - for i in range(len(mjm.geom_plugin)): - if mjm.geom_plugin[i] != -1: - p = mjm.geom_plugin[i] - geom_plugin_index[i] = len(plugin_id) - plugin_id.append(mjm.plugin[p]) - start = mjm.plugin_attradr[p] - end = mjm.plugin_attradr[p + 1] if p + 1 < mjm.nplugin else len(mjm.plugin_attr) - values = mjm.plugin_attr[start:end] - attr_values = [] - current = [] - for v in values: - if v == 0: - if current: - s = "".join(chr(int(x)) for x in current) - attr_values.append(float(s)) - current = [] - else: - current.append(v) - # Pad with zeros if less than 3 - attr_values += [0.0] * (3 - len(attr_values)) - plugin_attr.append(attr_values[:3]) - - plugin_id = np.array(plugin_id) - plugin_attr = np.array(plugin_attr) + # opt: check supported features in scalar flag types + for field, field_type in ( + (mjm.opt.disableflags, types.DisableBit), + (mjm.opt.enableflags, types.EnableBit), + ): + if field & ~np.bitwise_or.reduce(field_type): + raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.") if mjm.nflex > 1: raise NotImplementedError("Only one flex is unsupported.") @@ -137,15 +99,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any(): raise NotImplementedError("Flex collisions are not implemented.") - # check options - for opt, opt_types, msg in ( - (mjm.opt.integrator, types.IntegratorType, "Integrator"), - (mjm.opt.cone, types.ConeType, "Cone"), - (mjm.opt.solver, types.SolverType, "Solver"), - ): - if opt not in set(opt_types): - raise NotImplementedError(f"{msg} {opt} is unsupported.") - if mjm.opt.noslip_iterations > 0: raise NotImplementedError(f"noslip solver not implemented.") @@ -155,187 +108,116 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: ): raise NotImplementedError(f"Implicit integrators and fluid model not implemented.") + if (mjm.body_plugin != -1).any(): + raise NotImplementedError("Body plugins not supported.") + + if (mjm.actuator_plugin != -1).any(): + raise NotImplementedError("Actuator plugins not supported.") + + if (mjm.sensor_plugin != -1).any(): + raise NotImplementedError("Sensor plugins not supported.") + # 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: raise ValueError(f"Dense is unsupported for nv > {nv_max} (nv = {mjm.nv}).") - is_sparse = mujoco.mj_isSparse(mjm) + collision_sensors = (mujoco.mjtSensor.mjSENS_GEOMDIST, mujoco.mjtSensor.mjSENS_GEOMNORMAL, mujoco.mjtSensor.mjSENS_GEOMFROMTO) + is_collision_sensor = np.isin(mjm.sensor_type, collision_sensors) - # dof lower triangle row and column indices (used in solver) - dof_tri_row, dof_tri_col = np.tril_indices(mjm.nv) + def not_implemented(objtype, objid, geomtype): + if objtype == mujoco.mjtObj.mjOBJ_BODY: + geomnum = mjm.body_geomnum[objid] + geomadr = mjm.body_geomadr[objid] + for geomid in range(geomadr, geomadr + geomnum): + if mjm.geom_type[geomid] == geomtype: + return True + elif objtype == mujoco.mjtObj.mjOBJ_GEOM: + if mjm.geom_type[objid] == geomtype: + return True + return False - # indices for sparse qM_fullm (used in solver) - qM_fullm_i, qM_fullm_j = [], [] - for i in range(mjm.nv): - j = i - while j > -1: - qM_fullm_i.append(i) - qM_fullm_j.append(j) - j = mjm.dof_parentid[j] + for geoms in [ + (types.GeomType.BOX, types.GeomType.BOX), + (types.GeomType.CAPSULE, types.GeomType.BOX), + (types.GeomType.CYLINDER, types.GeomType.BOX), + (types.GeomType.PLANE, types.GeomType.BOX), + ]: + for objtype, objid, reftype, refid in zip( + mjm.sensor_objtype[is_collision_sensor], + mjm.sensor_objid[is_collision_sensor], + mjm.sensor_reftype[is_collision_sensor], + mjm.sensor_refid[is_collision_sensor], + ): + if not_implemented(objtype, objid, geoms[0]) and not_implemented(reftype, refid, geoms[1]): + raise NotImplementedError(f"Collision sensors with {geoms[0]} and {geoms[1]} are not implemented.") - # indices for sparse qM mul_m (used in support) - qM_mulm_i, qM_mulm_j, qM_madr_ij = [], [], [] - for i in range(mjm.nv): - madr_ij, j = mjm.dof_Madr[i], i + # create opt + opt = types.Option(**{f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)}) - while True: - madr_ij, j = madr_ij + 1, mjm.dof_parentid[j] - if j == -1: - break - qM_mulm_i.append(i) - qM_mulm_j.append(j) - qM_madr_ij.append(madr_ij) + # C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU + # adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly + # bouncing around the optimal solution + opt.tolerance = max(opt.tolerance, 1e-6) - # body_tree is a list of body ids grouped by tree level + # warp only fields + opt.is_sparse = bool(mujoco.mj_isSparse(mjm)) + ls_parallel_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "ls_parallel") + opt.ls_parallel = (ls_parallel_id > -1) and (mjm.numeric_data[mjm.numeric_adr[ls_parallel_id]] == 1) + opt.ls_parallel_min_step = 1.0e-6 # TODO(team): determine good default setting + opt.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0 + opt.broadphase = types.BroadphaseType.NXN + opt.broadphase_filter = types.BroadphaseFilter.PLANE | types.BroadphaseFilter.SPHERE | types.BroadphaseFilter.OBB + opt.graph_conditional = True + opt.run_collision_detection = True + opt.contact_sensor_maxmatch = 64 + + # place opt on device + for f in dataclasses.fields(types.Option): + if isinstance(f.type, wp.array): + setattr(opt, f.name, _create_array(getattr(opt, f.name), f.type, {"*": 1})) + else: + setattr(opt, f.name, f.type(getattr(opt, f.name))) + + # create stat + stat = types.Statistic(meaninertia=mjm.stat.meaninertia) + + # create model + m = types.Model(**{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model)}) + + m.opt = opt + m.stat = stat + + m.nacttrnbody = (mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY).sum() + m.nsensortaxel = mjm.mesh_vertnum[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_TACTILE]].sum() + m.nsensorcontact = (mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT).sum() + m.nrangefinder = (mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER).sum() + m.nmaxcondim = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max() + m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1)) + m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any() + m.block_dim = types.BlockDim() + + # body ids grouped by tree level bodies, body_depth = {}, np.zeros(mjm.nbody, dtype=int) - 1 for i in range(mjm.nbody): body_depth[i] = body_depth[mjm.body_parentid[i]] + 1 bodies.setdefault(body_depth[i], []).append(i) - body_tree = tuple(wp.array(bodies[i], dtype=int) for i in sorted(bodies)) + m.body_tree = tuple(wp.array(bodies[i], dtype=int) for i in sorted(bodies)) - # qLD_updates has dof tree ordering of qLD updates for sparse factor m - qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1 - - for k in range(mjm.nv): - # skip diagonal rows - if mjm.M_rownnz[k] == 1: - continue - dof_depth[k] = dof_depth[mjm.dof_parentid[k]] + 1 - i = mjm.dof_parentid[k] - diag_k = mjm.M_rowadr[k] + mjm.M_rownnz[k] - 1 - Madr_ki = diag_k - 1 - while i > -1: - qLD_updates.setdefault(dof_depth[i], []).append((i, k, Madr_ki)) - i = mjm.dof_parentid[i] - Madr_ki -= 1 - - qLD_updates = tuple(wp.array(qLD_updates[i], dtype=wp.vec3i) for i in sorted(qLD_updates)) - - # qM_tiles records the block diagonal structure of qM - tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1] - tiles = {} - for i in range(len(tile_corners)): - tile_beg = tile_corners[i] - tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1] - tiles.setdefault(tile_end - tile_beg, []).append(tile_beg) - - qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys())) - - # actuator_moment tiles are grouped by dof size and number of actuators - tree_id = np.arange(len(tile_corners), dtype=np.int32) - num_trees = int(np.max(tree_id)) if len(tree_id) > 0 else 0 - bodyid = [] - for i in range(mjm.nu): - trntype = mjm.actuator_trntype[i] - if trntype == mujoco.mjtTrn.mjTRN_JOINT or trntype == mujoco.mjtTrn.mjTRN_JOINTINPARENT: - jntid = mjm.actuator_trnid[i, 0] - bodyid.append(mjm.jnt_bodyid[jntid]) - elif trntype == mujoco.mjtTrn.mjTRN_TENDON: - tenid = mjm.actuator_trnid[i, 0] - adr = mjm.tendon_adr[tenid] - if mjm.wrap_type[adr] == mujoco.mjtWrap.mjWRAP_JOINT: - ten_num = mjm.tendon_num[tenid] - for i in range(ten_num): - bodyid.append(mjm.jnt_bodyid[mjm.wrap_objid[adr + i]]) - else: - for i in range(mjm.nv): - bodyid.append(mjm.dof_bodyid[i]) - elif trntype == mujoco.mjtTrn.mjTRN_BODY: - pass - elif trntype == mujoco.mjtTrn.mjTRN_SITE: - siteid = mjm.actuator_trnid[i, 0] - bid = mjm.site_bodyid[siteid] - while bid > 0: - bodyid.append(bid) - bid = mjm.body_parentid[bid] - elif trntype == mujoco.mjtTrn.mjTRN_SLIDERCRANK: - for i in range(mjm.nv): - bodyid.append(mjm.dof_bodyid[i]) - else: - raise NotImplementedError(f"Transmission type {trntype} not implemented.") - tree = mjm.body_treeid[np.array(bodyid, dtype=int)] - counts, ids = np.histogram(tree, bins=np.arange(0, num_trees + 2)) - acts_per_tree = dict(zip(ids, counts)) - - tiles = {} - act_beg = 0 - for i in range(len(tile_corners)): - tile_beg = tile_corners[i] - tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1] - tree = int(tree_id[i]) - act_num = acts_per_tree[tree] - tiles.setdefault((tile_end - tile_beg, act_num), []).append((tile_beg, act_beg)) - act_beg += act_num - - actuator_moment_tiles_nv, actuator_moment_tiles_nu = tuple(), tuple() - - for (nv, nu), adr in sorted(tiles.items()): - adr_nv = wp.array([nv for nv, _ in adr], dtype=int) - adr_nu = wp.array([nu for _, nu in adr], dtype=int) - actuator_moment_tiles_nv += (types.TileSet(adr=adr_nv, size=nv),) - actuator_moment_tiles_nu += (types.TileSet(adr=adr_nu, size=nu),) - - # fixed tendon - tendon_jnt_adr = [] - wrap_jnt_adr = [] - for i in range(mjm.ntendon): - adr = mjm.tendon_adr[i] - if mjm.wrap_type[adr] == mujoco.mjtWrap.mjWRAP_JOINT: - tendon_num = mjm.tendon_num[i] - for j in range(tendon_num): - tendon_jnt_adr.append(i) - wrap_jnt_adr.append(adr + j) - - # spatial tendon - tendon_site_pair_adr = [] - tendon_geom_adr = [] - - ten_wrapadr_site = [0] - ten_wrapnum_site = [] - for i, tendon_num in enumerate(mjm.tendon_num): - adr = mjm.tendon_adr[i] - # sites - if (mjm.wrap_type[adr : adr + tendon_num] == mujoco.mjtWrap.mjWRAP_SITE).all(): - if i < mjm.ntendon: - ten_wrapadr_site.append(ten_wrapadr_site[-1] + tendon_num) - ten_wrapnum_site.append(tendon_num) - else: - if i < mjm.ntendon: - ten_wrapadr_site.append(ten_wrapadr_site[-1]) - ten_wrapnum_site.append(0) - - # geoms - for j in range(tendon_num): - wrap_type = mjm.wrap_type[adr + j] - if j < tendon_num - 1: - next_wrap_type = mjm.wrap_type[adr + j + 1] - if wrap_type == mujoco.mjtWrap.mjWRAP_SITE and next_wrap_type == mujoco.mjtWrap.mjWRAP_SITE: - tendon_site_pair_adr.append(i) - if wrap_type == mujoco.mjtWrap.mjWRAP_SPHERE or wrap_type == mujoco.mjtWrap.mjWRAP_CYLINDER: - tendon_geom_adr.append(i) - - wrap_site_adr = np.nonzero(mjm.wrap_type == mujoco.mjtWrap.mjWRAP_SITE)[0] - wrap_site_pair_adr = np.setdiff1d(wrap_site_adr[np.nonzero(np.diff(wrap_site_adr) == 1)[0]], mjm.tendon_adr[1:] - 1) - wrap_geom_adr = np.nonzero(np.isin(mjm.wrap_type, [mujoco.mjtWrap.mjWRAP_SPHERE, mujoco.mjtWrap.mjWRAP_CYLINDER]))[0] - - # pulley scaling - wrap_pulley_scale = np.ones(mjm.nwrap, dtype=float) - pulley_adr = np.nonzero(mjm.wrap_type == mujoco.mjtWrap.mjWRAP_PULLEY)[0] - for tadr, tnum in zip(mjm.tendon_adr, mjm.tendon_num): - for padr in pulley_adr: - if tadr <= padr < tadr + tnum: - wrap_pulley_scale[padr : tadr + tnum] = 1.0 / mjm.wrap_prm[padr] - - # mocap - mocap_bodyid = np.arange(mjm.nbody)[mjm.body_mocapid >= 0] - mocap_bodyid = mocap_bodyid[mjm.body_mocapid[mjm.body_mocapid >= 0].argsort()] + m.mocap_bodyid = np.arange(mjm.nbody)[mjm.body_mocapid >= 0] + m.mocap_bodyid = m.mocap_bodyid[mjm.body_mocapid[mjm.body_mocapid >= 0].argsort()] + m.body_fluid_ellipsoid = np.zeros(mjm.nbody, dtype=bool) + m.body_fluid_ellipsoid[mjm.geom_bodyid[mjm.geom_fluid.reshape(mjm.ngeom, mujoco.mjNFLUID)[:, 0] > 0]] = True + jnt_limited_slide_hinge = mjm.jnt_limited & np.isin(mjm.jnt_type, (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE)) + m.jnt_limited_slide_hinge_adr = np.nonzero(jnt_limited_slide_hinge)[0] + m.jnt_limited_ball_adr = np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0] + m.dof_tri_row, m.dof_tri_col = np.tril_indices(mjm.nv) # precalculated geom pairs filterparent = not (mjm.opt.disableflags & types.DisableBit.FILTERPARENT) geom1, geom2 = np.triu_indices(mjm.ngeom, k=1) - nxn_geom_pair = np.stack((geom1, geom2), axis=1) + m.nxn_geom_pair = np.stack((geom1, geom2), axis=1) bodyid1 = mjm.geom_bodyid[geom1] bodyid2 = mjm.geom_bodyid[geom2] @@ -359,83 +241,18 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: nxn_pairid_contact[~(mask & ~self_collision & ~parent_child_collision & ~exclude)] = -2 # contact pairs + def upper_tri_index(n, i, j): + i, j = (j, i) if j < i else (i, j) + return (i * (2 * n - i - 3)) // 2 + j - 1 + for i in range(mjm.npair): - pair_geom1 = mjm.pair_geom1[i] - pair_geom2 = mjm.pair_geom2[i] + nxn_pairid_contact[upper_tri_index(mjm.ngeom, mjm.pair_geom1[i], mjm.pair_geom2[i])] = i - if pair_geom2 < pair_geom1: - pairid = np.int32(math.upper_tri_index(mjm.ngeom, int(pair_geom2), int(pair_geom1))) - else: - pairid = np.int32(math.upper_tri_index(mjm.ngeom, int(pair_geom1), int(pair_geom2))) - - nxn_pairid_contact[pairid] = i - - def create_nmodel_batched_array(mjm_array, dtype, expand_dim=True): - array = wp.array(mjm_array, dtype=dtype) - # add private attribute for JAX to determine which fields are batched - array._is_batched = True - if not expand_dim: - array.strides = (0,) + array.strides[1:] - array.shape = (1,) + array.shape[1:] - return array - array.strides = (0,) + array.strides - array.ndim += 1 - array.shape = (1,) + array.shape - return array - - # rangefinder - is_rangefinder = mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER - sensor_rangefinder_adr = np.nonzero(is_rangefinder)[0] - rangefinder_sensor_adr = np.full(mjm.nsensor, -1) - rangefinder_sensor_adr[sensor_rangefinder_adr] = np.arange(len(sensor_rangefinder_adr)) - - # contact sensor - sensor_adr_to_contact_adr = np.clip(np.cumsum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT) - 1, a_min=0, a_max=None) - - condim = np.concatenate((mjm.geom_condim, mjm.pair_dim)) - condim_max = np.max(condim) if len(condim) > 0 else 0 - - # collision sensors - is_collision_sensor = np.isin( - mjm.sensor_type, [mujoco.mjtSensor.mjSENS_GEOMDIST, mujoco.mjtSensor.mjSENS_GEOMNORMAL, mujoco.mjtSensor.mjSENS_GEOMFROMTO] - ) sensor_collision_adr = np.nonzero(is_collision_sensor)[0] - collision_sensor_adr = np.full(mjm.nsensor, -1) - collision_sensor_adr[sensor_collision_adr] = np.arange(len(sensor_collision_adr)) - - if is_collision_sensor.any(): - - def not_implemented(objtype, objid, geomtype): - if objtype == mujoco.mjtObj.mjOBJ_BODY: - geomnum = mjm.body_geomnum[objid] - geomadr = mjm.body_geomadr[objid] - for geomid in range(geomadr, geomadr + geomnum): - if mjm.geom_type[geomid] == geomtype: - return True - elif objtype == mujoco.mjtObj.mjOBJ_GEOM: - if mjm.geom_type[objid] == geomtype: - return True - return False - - for geoms in [ - (types.GeomType.BOX, types.GeomType.BOX), - (types.GeomType.CAPSULE, types.GeomType.BOX), - (types.GeomType.CYLINDER, types.GeomType.BOX), - (types.GeomType.PLANE, types.GeomType.BOX), - ]: - for objtype, objid, reftype, refid in zip( - mjm.sensor_objtype[is_collision_sensor], - mjm.sensor_objid[is_collision_sensor], - mjm.sensor_reftype[is_collision_sensor], - mjm.sensor_refid[is_collision_sensor], - ): - if not_implemented(objtype, objid, geoms[0]) and not_implemented(reftype, refid, geoms[1]): - raise NotImplementedError(f"Collision sensors with {geoms[0]} and {geoms[1]} are not implemented.") - nxn_pairid_collision = -1 * np.ones(len(geom1), dtype=int) pairids = [] collision_geom_adr = [0] - sensor_collision_start_adr = [] + m.sensor_collision_start_adr = [] for i in range(sensor_collision_adr.size): sensorid = sensor_collision_adr[i] objtype = mjm.sensor_objtype[sensorid] @@ -461,515 +278,260 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: geomid = 0 for geom1id in range(id1, id1 + n1): for geom2id in range(id2, id2 + n2): - if geom2id < geom1id: - pairid = np.int32(math.upper_tri_index(mjm.ngeom, int(geom2id), int(geom1id))) - else: - pairid = np.int32(math.upper_tri_index(mjm.ngeom, int(geom1id), int(geom2id))) + pairid = upper_tri_index(mjm.ngeom, geom1id, geom2id) if pairid in pairids: - sensor_collision_start_adr.append(nxn_pairid_collision[pairid]) + m.sensor_collision_start_adr.append(nxn_pairid_collision[pairid]) else: pairids.append(pairid) adr = collision_geom_adr[-1] + geomid nxn_pairid_collision[pairid] = adr - sensor_collision_start_adr.append(adr) + m.sensor_collision_start_adr.append(adr) geomid += 1 if i < sensor_collision_adr.size - 1: collision_geom_adr.append(collision_geom_adr[-1] + n1 * n2) - include = (nxn_pairid_contact > -2) | (nxn_pairid_collision >= 0) - nxn_pairid = np.hstack([nxn_pairid_contact.reshape((-1, 1)), nxn_pairid_collision.reshape((-1, 1))]) - nxn_pairid_filtered = nxn_pairid[include] - nxn_geom_pair_filtered = nxn_geom_pair[include] + m.nsensorcollision = (nxn_pairid_collision >= 0).sum() + nxn_include = (nxn_pairid_contact > -2) | (nxn_pairid_collision >= 0) + + if nxn_include.sum() < 250_000: + opt.broadphase = types.BroadphaseType.NXN + elif mjm.ngeom < 1000: + opt.broadphase = types.BroadphaseType.SAP_TILE + else: + opt.broadphase = types.BroadphaseType.SAP_SEGMENTED + + m.nxn_geom_pair_filtered = m.nxn_geom_pair[nxn_include] + m.nxn_pairid = np.hstack([nxn_pairid_contact.reshape((-1, 1)), nxn_pairid_collision.reshape((-1, 1))]) + m.nxn_pairid_filtered = m.nxn_pairid[nxn_include] # count contact pair types - geom_type_pair_count = np.bincount( + def geom_trid_index(i, j): + i, j = (j, i) if j < i else (i, j) + return (i * (2 * len(types.GeomType) - i - 1)) // 2 + j + + m.geom_pair_type_count = tuple( + np.bincount( + [geom_trid_index(mjm.geom_type[geom1[i]], mjm.geom_type[geom2[i]]) for i in np.arange(len(geom1)) if nxn_include[i]], + minlength=len(types.GeomType) * (len(types.GeomType) + 1) // 2, + ) + ) + + # compute nmaxpolygon and nmaxmeshdeg given the geom pairs for the model + nboxbox = m.geom_pair_type_count[geom_trid_index(types.GeomType.BOX, types.GeomType.BOX)] + nboxmesh = m.geom_pair_type_count[geom_trid_index(types.GeomType.BOX, types.GeomType.MESH)] + nmeshmesh = m.geom_pair_type_count[geom_trid_index(types.GeomType.MESH, types.GeomType.MESH)] + # need at least 4 (square sides) if there's a box collision needing multiccd + m.nmaxpolygon = 4 * (nboxbox + nboxmesh > 0) + m.nmaxmeshdeg = 3 * (nboxbox + nboxmesh > 0) + # possibly need to allocate more memory if there's meshes + if nmeshmesh + nboxmesh > 0: + # TODO(kbayes): remove nboxbox or enable ccd for box-box collisions + m.nmaxpolygon = np.append(mjm.mesh_polyvertnum, m.nmaxpolygon).max() + m.nmaxmeshdeg = np.append(mjm.mesh_polymapnum, m.nmaxmeshdeg).max() + + # filter plugins for only geom plugins, drop the rest + m.plugin, m.plugin_attr = [], [] + m.geom_plugin_index = np.full_like(mjm.geom_type, -1) + + for i in range(len(mjm.geom_plugin)): + if mjm.geom_plugin[i] == -1: + continue + p = mjm.geom_plugin[i] + m.geom_plugin_index[i] = len(m.plugin) + m.plugin.append(mjm.plugin[p]) + start = mjm.plugin_attradr[p] + end = mjm.plugin_attradr[p + 1] if p + 1 < mjm.nplugin else len(mjm.plugin_attr) + values = mjm.plugin_attr[start:end] + attr_values = [] + current = [] + for v in values: + if v == 0: + if current: + s = "".join(chr(int(x)) for x in current) + attr_values.append(float(s)) + current = [] + else: + current.append(v) + # Pad with zeros if less than 3 + attr_values += [0.0] * (3 - len(attr_values)) + m.plugin_attr.append(attr_values[:3]) + + # equality constraint addresses + m.eq_connect_adr = np.nonzero(mjm.eq_type == types.EqType.CONNECT)[0] + m.eq_wld_adr = np.nonzero(mjm.eq_type == types.EqType.WELD)[0] + m.eq_jnt_adr = np.nonzero(mjm.eq_type == types.EqType.JOINT)[0] + m.eq_ten_adr = np.nonzero(mjm.eq_type == types.EqType.TENDON)[0] + + # fixed tendon + m.tendon_jnt_adr, m.wrap_jnt_adr = [], [] + for i in range(mjm.ntendon): + adr = mjm.tendon_adr[i] + if mjm.wrap_type[adr] == mujoco.mjtWrap.mjWRAP_JOINT: + tendon_num = mjm.tendon_num[i] + for j in range(tendon_num): + m.tendon_jnt_adr.append(i) + m.wrap_jnt_adr.append(adr + j) + + # spatial tendon + m.tendon_site_pair_adr, m.tendon_geom_adr = [], [] + m.ten_wrapadr_site, m.ten_wrapnum_site = [0], [] + for i, tendon_num in enumerate(mjm.tendon_num): + adr = mjm.tendon_adr[i] + # sites + if (mjm.wrap_type[adr : adr + tendon_num] == mujoco.mjtWrap.mjWRAP_SITE).all(): + if i < mjm.ntendon: + m.ten_wrapadr_site.append(m.ten_wrapadr_site[-1] + tendon_num) + m.ten_wrapnum_site.append(tendon_num) + else: + if i < mjm.ntendon: + m.ten_wrapadr_site.append(m.ten_wrapadr_site[-1]) + m.ten_wrapnum_site.append(0) + + # geoms + for j in range(tendon_num): + wrap_type = mjm.wrap_type[adr + j] + if j < tendon_num - 1: + next_wrap_type = mjm.wrap_type[adr + j + 1] + if wrap_type == mujoco.mjtWrap.mjWRAP_SITE and next_wrap_type == mujoco.mjtWrap.mjWRAP_SITE: + m.tendon_site_pair_adr.append(i) + if wrap_type == mujoco.mjtWrap.mjWRAP_SPHERE or wrap_type == mujoco.mjtWrap.mjWRAP_CYLINDER: + m.tendon_geom_adr.append(i) + + m.tendon_limited_adr = np.nonzero(mjm.tendon_limited)[0] + m.wrap_site_adr = np.nonzero(mjm.wrap_type == mujoco.mjtWrap.mjWRAP_SITE)[0] + m.wrap_site_pair_adr = np.setdiff1d(m.wrap_site_adr[np.nonzero(np.diff(m.wrap_site_adr) == 1)[0]], mjm.tendon_adr[1:] - 1) + m.wrap_geom_adr = np.nonzero(np.isin(mjm.wrap_type, [mujoco.mjtWrap.mjWRAP_SPHERE, mujoco.mjtWrap.mjWRAP_CYLINDER]))[0] + + # pulley scaling + m.wrap_pulley_scale = np.ones(mjm.nwrap, dtype=float) + pulley_adr = np.nonzero(mjm.wrap_type == mujoco.mjtWrap.mjWRAP_PULLEY)[0] + for tadr, tnum in zip(mjm.tendon_adr, mjm.tendon_num): + for padr in pulley_adr: + if tadr <= padr < tadr + tnum: + m.wrap_pulley_scale[padr : tadr + tnum] = 1.0 / mjm.wrap_prm[padr] + + m.actuator_trntype_body_adr = np.nonzero(mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY)[0] + + # sensor addresses + m.sensor_pos_adr = np.nonzero( + (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS) + & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_JOINTLIMITPOS) + & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONLIMITPOS) + )[0] + m.sensor_limitpos_adr = np.nonzero( + (mjm.sensor_type == mujoco.mjtSensor.mjSENS_JOINTLIMITPOS) | (mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONLIMITPOS) + )[0] + m.sensor_vel_adr = np.nonzero( + (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL) + & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_JOINTLIMITVEL) + & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONLIMITVEL) + )[0] + m.sensor_limitvel_adr = np.nonzero( + (mjm.sensor_type == mujoco.mjtSensor.mjSENS_JOINTLIMITVEL) | (mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONLIMITVEL) + )[0] + m.sensor_acc_adr = np.nonzero( + (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC) + & ( + (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TOUCH) + | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_JOINTLIMITFRC) + | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONLIMITFRC) + | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONACTFRC) + ) + )[0] + m.sensor_rangefinder_adr = np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER)[0] + m.rangefinder_sensor_adr = np.full(mjm.nsensor, -1) + m.rangefinder_sensor_adr[m.sensor_rangefinder_adr] = np.arange(len(m.sensor_rangefinder_adr)) + m.collision_sensor_adr = np.full(mjm.nsensor, -1) + m.collision_sensor_adr[sensor_collision_adr] = np.arange(len(sensor_collision_adr)) + m.sensor_touch_adr = np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TOUCH)[0] + limitfrc_sensors = (mujoco.mjtSensor.mjSENS_JOINTLIMITFRC, mujoco.mjtSensor.mjSENS_TENDONLIMITFRC) + m.sensor_limitfrc_adr = np.nonzero(np.isin(mjm.sensor_type, limitfrc_sensors))[0] + m.sensor_e_potential = (mjm.sensor_type == mujoco.mjtSensor.mjSENS_E_POTENTIAL).any() + m.sensor_e_kinetic = (mjm.sensor_type == mujoco.mjtSensor.mjSENS_E_KINETIC).any() + m.sensor_tendonactfrc_adr = np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONACTFRC)[0] + subtreevel_sensors = (mujoco.mjtSensor.mjSENS_SUBTREELINVEL, mujoco.mjtSensor.mjSENS_SUBTREEANGMOM) + m.sensor_subtree_vel = np.isin(mjm.sensor_type, subtreevel_sensors).any() + m.sensor_contact_adr = np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT)[0] + m.sensor_adr_to_contact_adr = np.clip(np.cumsum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT) - 1, a_min=0, a_max=None) + m.sensor_rne_postconstraint = np.isin( + mjm.sensor_type, [ - math.upper_trid_index(len(types.GeomType), int(mjm.geom_type[geom1[i]]), int(mjm.geom_type[geom2[i]])) - for i in np.arange(len(geom1)) - if nxn_pairid_contact[i] > -2 or nxn_pairid_collision[i] > -1 + mujoco.mjtSensor.mjSENS_ACCELEROMETER, + mujoco.mjtSensor.mjSENS_FORCE, + mujoco.mjtSensor.mjSENS_TORQUE, + mujoco.mjtSensor.mjSENS_FRAMELINACC, + mujoco.mjtSensor.mjSENS_FRAMEANGACC, ], - minlength=len(types.GeomType) * (len(types.GeomType) + 1) // 2, - ) + ).any() + m.sensor_rangefinder_bodyid = mjm.site_bodyid[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER]] + m.taxel_vertadr = [ + 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]]) + ] + m.taxel_sensorid = [ + 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]]) + ] - # disable collisions if there are no potentially colliding pairs - disableflags = mjm.opt.disableflags - if np.sum(geom_type_pair_count) == 0: - disableflags |= types.DisableBit.CONTACT + # qM_tiles records the block diagonal structure of qM + tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1] + tiles = {} + for i in range(len(tile_corners)): + tile_beg = tile_corners[i] + tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1] + tiles.setdefault(tile_end - tile_beg, []).append(tile_beg) + m.qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys())) - if mjm.geom_fluid.size: - geom_fluid_params = mjm.geom_fluid.reshape(mjm.ngeom, mujoco.mjNFLUID) - else: - geom_fluid_params = np.zeros((mjm.ngeom, mujoco.mjNFLUID)) + # qLD_updates has dof tree ordering of qLD updates for sparse factor m + qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1 - body_fluid_ellipsoid = np.zeros(mjm.nbody, dtype=bool) - if mjm.ngeom: - active_geom = geom_fluid_params[:, 0] > 0 - if np.any(active_geom): - body_fluid_ellipsoid[mjm.geom_bodyid[active_geom]] = True + for k in range(mjm.nv): + # skip diagonal rows + if mjm.M_rownnz[k] == 1: + continue + dof_depth[k] = dof_depth[mjm.dof_parentid[k]] + 1 + i = mjm.dof_parentid[k] + diag_k = mjm.M_rowadr[k] + mjm.M_rownnz[k] - 1 + Madr_ki = diag_k - 1 + while i > -1: + qLD_updates.setdefault(dof_depth[i], []).append((i, k, Madr_ki)) + i = mjm.dof_parentid[i] + Madr_ki -= 1 + m.qLD_updates = tuple(wp.array(qLD_updates[i], dtype=wp.vec3i) for i in sorted(qLD_updates)) - if nxn_geom_pair_filtered.shape[0] < 250_000: - broadphase = types.BroadphaseType.NXN - elif mjm.ngeom < 1000: - broadphase = types.BroadphaseType.SAP_TILE - else: - broadphase = types.BroadphaseType.SAP_SEGMENTED + # indices for sparse qM_fullm (used in solver) + m.qM_fullm_i, m.qM_fullm_j = [], [] + for i in range(mjm.nv): + j = i + while j > -1: + m.qM_fullm_i.append(i) + m.qM_fullm_j.append(j) + j = mjm.dof_parentid[j] - ls_parallel_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "ls_parallel") - if (ls_parallel_id > -1) and (mjm.numeric_data[mjm.numeric_adr[ls_parallel_id]] == 1): - ls_parallel = True - else: - ls_parallel = False + # indices for sparse qM mul_m (used in support) + m.qM_mulm_i, m.qM_mulm_j, m.qM_madr_ij = [], [], [] + for i in range(mjm.nv): + madr_ij, j = mjm.dof_Madr[i], i - m = types.Model( - nq=mjm.nq, - nv=mjm.nv, - nu=mjm.nu, - na=mjm.na, - nbody=mjm.nbody, - njnt=mjm.njnt, - nM=mjm.nM, - nC=mjm.nC, - ngeom=mjm.ngeom, - nsite=mjm.nsite, - ncam=mjm.ncam, - nlight=mjm.nlight, - nflex=mjm.nflex, - nflexvert=mjm.nflexvert, - nflexedge=mjm.nflexedge, - nflexelem=mjm.nflexelem, - nflexelemdata=mjm.nflexelemdata, - nmeshvert=mjm.nmeshvert, - nmeshface=mjm.nmeshface, - nmeshgraph=mjm.nmeshgraph, - nmeshpoly=mjm.nmeshpoly, - nmeshpolyvert=mjm.nmeshpolyvert, - nmeshpolymap=mjm.nmeshpolymap, - nhfield=mjm.nhfield, - nhfielddata=mjm.nhfielddata, - nmat=mjm.nmat, - npair=mjm.npair, - nexclude=mjm.nexclude, - neq=mjm.neq, - ntendon=mjm.ntendon, - nwrap=mjm.nwrap, - nsensor=mjm.nsensor, - nmocap=mjm.nmocap, - ngravcomp=mjm.ngravcomp, - nsensordata=mjm.nsensordata, - opt=types.Option( - timestep=create_nmodel_batched_array(np.array(mjm.opt.timestep), dtype=float, expand_dim=False), - tolerance=create_nmodel_batched_array( - np.array(np.maximum(mjm.opt.tolerance, _TOLERANCE_F32)), dtype=float, expand_dim=False - ), - ls_tolerance=create_nmodel_batched_array(np.array(mjm.opt.ls_tolerance), dtype=float, expand_dim=False), - ccd_tolerance=create_nmodel_batched_array(np.array(mjm.opt.ccd_tolerance), dtype=float, expand_dim=False), - gravity=create_nmodel_batched_array(mjm.opt.gravity, dtype=wp.vec3, expand_dim=False), - magnetic=create_nmodel_batched_array(mjm.opt.magnetic, dtype=wp.vec3, expand_dim=False), - wind=create_nmodel_batched_array(mjm.opt.wind, dtype=wp.vec3, expand_dim=False), - has_fluid=bool(mjm.opt.wind.any() or mjm.opt.density or mjm.opt.viscosity), - density=create_nmodel_batched_array(np.array(mjm.opt.density), dtype=float, expand_dim=False), - viscosity=create_nmodel_batched_array(np.array(mjm.opt.viscosity), dtype=float, expand_dim=False), - cone=mjm.opt.cone, - solver=mjm.opt.solver, - iterations=mjm.opt.iterations, - ls_iterations=mjm.opt.ls_iterations, - integrator=mjm.opt.integrator, - disableflags=type(mjm.opt.disableflags)(disableflags), - enableflags=mjm.opt.enableflags, - impratio=create_nmodel_batched_array(np.array(mjm.opt.impratio), dtype=float, expand_dim=False), - is_sparse=bool(is_sparse), - ls_parallel=ls_parallel, - ls_parallel_min_step=1.0e-6, # TODO(team): determine good default setting - ccd_iterations=mjm.opt.ccd_iterations, - broadphase=int(broadphase), - broadphase_filter=int(types.BroadphaseFilter.PLANE | types.BroadphaseFilter.SPHERE | types.BroadphaseFilter.OBB), - graph_conditional=True, - sdf_initpoints=mjm.opt.sdf_initpoints, - sdf_iterations=mjm.opt.sdf_iterations, - run_collision_detection=True, - contact_sensor_maxmatch=64, - ), - stat=types.Statistic( - meaninertia=mjm.stat.meaninertia, - ), - qpos0=create_nmodel_batched_array(mjm.qpos0, dtype=float), - qpos_spring=create_nmodel_batched_array(mjm.qpos_spring, dtype=float), - body_parentid=wp.array(mjm.body_parentid, dtype=int), - body_rootid=wp.array(mjm.body_rootid, dtype=int), - body_weldid=wp.array(mjm.body_weldid, dtype=int), - body_mocapid=wp.array(mjm.body_mocapid, dtype=int), - body_jntnum=wp.array(mjm.body_jntnum, dtype=int), - body_jntadr=wp.array(mjm.body_jntadr, dtype=int), - body_dofnum=wp.array(mjm.body_dofnum, dtype=int), - body_dofadr=wp.array(mjm.body_dofadr, dtype=int), - body_geomnum=wp.array(mjm.body_geomnum, dtype=int), - body_geomadr=wp.array(mjm.body_geomadr, dtype=int), - body_pos=create_nmodel_batched_array(mjm.body_pos, dtype=wp.vec3), - body_quat=create_nmodel_batched_array(mjm.body_quat, dtype=wp.quat), - body_ipos=create_nmodel_batched_array(mjm.body_ipos, dtype=wp.vec3), - body_iquat=create_nmodel_batched_array(mjm.body_iquat, dtype=wp.quat), - body_mass=create_nmodel_batched_array(mjm.body_mass, dtype=float), - body_subtreemass=create_nmodel_batched_array(mjm.body_subtreemass, dtype=float), - body_inertia=create_nmodel_batched_array(mjm.body_inertia, dtype=wp.vec3), - body_invweight0=create_nmodel_batched_array(mjm.body_invweight0, dtype=wp.vec2), - body_gravcomp=create_nmodel_batched_array(mjm.body_gravcomp, dtype=float), - body_contype=wp.array(mjm.body_contype, dtype=int), - body_conaffinity=wp.array(mjm.body_conaffinity, dtype=int), - oct_aabb=wp.array2d(mjm.oct_aabb, dtype=wp.vec3), - oct_child=wp.array(mjm.oct_child, dtype=types.vec8i), - oct_coeff=wp.array(mjm.oct_coeff, dtype=types.vec8f), - jnt_type=wp.array(mjm.jnt_type, dtype=int), - jnt_qposadr=wp.array(mjm.jnt_qposadr, dtype=int), - jnt_dofadr=wp.array(mjm.jnt_dofadr, dtype=int), - jnt_bodyid=wp.array(mjm.jnt_bodyid, dtype=int), - jnt_limited=wp.array(mjm.jnt_limited, dtype=int), - jnt_actfrclimited=wp.array(mjm.jnt_actfrclimited, dtype=bool), - jnt_actgravcomp=wp.array(mjm.jnt_actgravcomp, dtype=int), - jnt_solref=create_nmodel_batched_array(mjm.jnt_solref, dtype=wp.vec2), - jnt_solimp=create_nmodel_batched_array(mjm.jnt_solimp, dtype=types.vec5), - jnt_pos=create_nmodel_batched_array(mjm.jnt_pos, dtype=wp.vec3), - jnt_axis=create_nmodel_batched_array(mjm.jnt_axis, dtype=wp.vec3), - jnt_stiffness=create_nmodel_batched_array(mjm.jnt_stiffness, dtype=float), - jnt_range=create_nmodel_batched_array(mjm.jnt_range, dtype=wp.vec2), - jnt_actfrcrange=create_nmodel_batched_array(mjm.jnt_actfrcrange, dtype=wp.vec2), - jnt_margin=create_nmodel_batched_array(mjm.jnt_margin, dtype=float), - dof_bodyid=wp.array(mjm.dof_bodyid, dtype=int), - dof_jntid=wp.array(mjm.dof_jntid, dtype=int), - dof_parentid=wp.array(mjm.dof_parentid, dtype=int), - dof_Madr=wp.array(mjm.dof_Madr, dtype=int), - dof_armature=create_nmodel_batched_array(mjm.dof_armature, dtype=float), - dof_damping=create_nmodel_batched_array(mjm.dof_damping, dtype=float), - dof_invweight0=create_nmodel_batched_array(mjm.dof_invweight0, dtype=float), - dof_frictionloss=create_nmodel_batched_array(mjm.dof_frictionloss, dtype=float), - dof_solimp=create_nmodel_batched_array(mjm.dof_solimp, dtype=types.vec5), - dof_solref=create_nmodel_batched_array(mjm.dof_solref, dtype=wp.vec2), - geom_type=wp.array(mjm.geom_type, dtype=int), - geom_contype=wp.array(mjm.geom_contype, dtype=int), - geom_conaffinity=wp.array(mjm.geom_conaffinity, dtype=int), - geom_condim=wp.array(mjm.geom_condim, dtype=int), - geom_bodyid=wp.array(mjm.geom_bodyid, dtype=int), - geom_dataid=wp.array(mjm.geom_dataid, dtype=int), - geom_matid=create_nmodel_batched_array(mjm.geom_matid, dtype=int), - geom_group=wp.array(mjm.geom_group, dtype=int), - geom_priority=wp.array(mjm.geom_priority, dtype=int), - geom_solmix=create_nmodel_batched_array(mjm.geom_solmix, dtype=float), - geom_solref=create_nmodel_batched_array(mjm.geom_solref, dtype=wp.vec2), - geom_solimp=create_nmodel_batched_array(mjm.geom_solimp, dtype=types.vec5), - geom_size=create_nmodel_batched_array(mjm.geom_size, dtype=wp.vec3), - geom_aabb=create_nmodel_batched_array(mjm.geom_aabb, dtype=wp.vec3), - geom_rbound=create_nmodel_batched_array(mjm.geom_rbound, dtype=float), - geom_pos=create_nmodel_batched_array(mjm.geom_pos, dtype=wp.vec3), - geom_quat=create_nmodel_batched_array(mjm.geom_quat, dtype=wp.quat), - geom_friction=create_nmodel_batched_array(mjm.geom_friction, dtype=wp.vec3), - geom_margin=create_nmodel_batched_array(mjm.geom_margin, dtype=float), - geom_gap=create_nmodel_batched_array(mjm.geom_gap, dtype=float), - geom_fluid=wp.array(geom_fluid_params, dtype=float), - geom_rgba=create_nmodel_batched_array(mjm.geom_rgba, dtype=wp.vec4), - site_type=wp.array(mjm.site_type, dtype=int), - site_bodyid=wp.array(mjm.site_bodyid, dtype=int), - site_size=wp.array(mjm.site_size, dtype=wp.vec3), - site_pos=create_nmodel_batched_array(mjm.site_pos, dtype=wp.vec3), - site_quat=create_nmodel_batched_array(mjm.site_quat, dtype=wp.quat), - cam_mode=wp.array(mjm.cam_mode, dtype=int), - cam_bodyid=wp.array(mjm.cam_bodyid, dtype=int), - cam_targetbodyid=wp.array(mjm.cam_targetbodyid, dtype=int), - cam_pos=create_nmodel_batched_array(mjm.cam_pos, dtype=wp.vec3), - cam_quat=create_nmodel_batched_array(mjm.cam_quat, dtype=wp.quat), - cam_poscom0=create_nmodel_batched_array(mjm.cam_poscom0, dtype=wp.vec3), - cam_pos0=create_nmodel_batched_array(mjm.cam_pos0, dtype=wp.vec3), - cam_mat0=create_nmodel_batched_array(mjm.cam_mat0, dtype=wp.mat33), - cam_fovy=wp.array(mjm.cam_fovy, dtype=float), - cam_resolution=wp.array(mjm.cam_resolution, dtype=wp.vec2i), - cam_sensorsize=wp.array(mjm.cam_sensorsize, dtype=wp.vec2), - cam_intrinsic=wp.array(mjm.cam_intrinsic, dtype=wp.vec4), - light_mode=wp.array(mjm.light_mode, dtype=int), - light_bodyid=wp.array(mjm.light_bodyid, dtype=int), - light_targetbodyid=wp.array(mjm.light_targetbodyid, dtype=int), - light_type=create_nmodel_batched_array(mjm.light_type, dtype=int), - light_castshadow=create_nmodel_batched_array(mjm.light_castshadow, dtype=bool), - light_active=create_nmodel_batched_array(mjm.light_active, dtype=bool), - light_pos=create_nmodel_batched_array(mjm.light_pos, dtype=wp.vec3), - light_dir=create_nmodel_batched_array(mjm.light_dir, dtype=wp.vec3), - light_poscom0=create_nmodel_batched_array(mjm.light_poscom0, dtype=wp.vec3), - light_pos0=create_nmodel_batched_array(mjm.light_pos0, dtype=wp.vec3), - light_dir0=create_nmodel_batched_array(mjm.light_dir0, dtype=wp.vec3), - flex_dim=wp.array(mjm.flex_dim, dtype=int), - flex_vertadr=wp.array(mjm.flex_vertadr, dtype=int), - flex_vertnum=wp.array(mjm.flex_vertnum, dtype=int), - flex_edgeadr=wp.array(mjm.flex_edgeadr, dtype=int), - flex_elemedgeadr=wp.array(mjm.flex_elemedgeadr, dtype=int), - flex_vertbodyid=wp.array(mjm.flex_vertbodyid, dtype=int), - flex_edge=wp.array(mjm.flex_edge, dtype=wp.vec2i), - flex_edgeflap=wp.array(mjm.flex_edgeflap, dtype=wp.vec2i), - flex_elem=wp.array(mjm.flex_elem, dtype=int), - flex_elemedge=wp.array(mjm.flex_elemedge, dtype=int), - flexedge_length0=wp.array(mjm.flexedge_length0, dtype=float), - flex_stiffness=wp.array(mjm.flex_stiffness.flatten(), dtype=float), - flex_bending=wp.array(mjm.flex_bending.flatten(), dtype=float), - flex_damping=wp.array(mjm.flex_damping, dtype=float), - mesh_vertadr=wp.array(mjm.mesh_vertadr, dtype=int), - mesh_vertnum=wp.array(mjm.mesh_vertnum, dtype=int), - mesh_faceadr=wp.array(mjm.mesh_faceadr, dtype=int), - mesh_normaladr=wp.array(mjm.mesh_normaladr, dtype=int), - mesh_graphadr=wp.array(mjm.mesh_graphadr, dtype=int), - mesh_vert=wp.array(mjm.mesh_vert, dtype=wp.vec3), - mesh_normal=wp.array(mjm.mesh_normal, dtype=wp.vec3), - mesh_face=wp.array(mjm.mesh_face, dtype=wp.vec3i), - 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), - mesh_polyvertadr=wp.array(mjm.mesh_polyvertadr, dtype=int), - mesh_polyvertnum=wp.array(mjm.mesh_polyvertnum, dtype=int), - mesh_polyvert=wp.array(mjm.mesh_polyvert, dtype=int), - mesh_polymapadr=wp.array(mjm.mesh_polymapadr, dtype=int), - mesh_polymapnum=wp.array(mjm.mesh_polymapnum, dtype=int), - mesh_polymap=wp.array(mjm.mesh_polymap, dtype=int), - hfield_size=wp.array(mjm.hfield_size, dtype=wp.vec4), - hfield_nrow=wp.array(mjm.hfield_nrow, dtype=int), - hfield_ncol=wp.array(mjm.hfield_ncol, dtype=int), - hfield_adr=wp.array(mjm.hfield_adr, dtype=int), - hfield_data=wp.array(mjm.hfield_data, dtype=float), - mat_texid=create_nmodel_batched_array(mjm.mat_texid, dtype=int), - mat_texrepeat=create_nmodel_batched_array(mjm.mat_texrepeat, dtype=wp.vec2), - mat_rgba=create_nmodel_batched_array(mjm.mat_rgba, dtype=wp.vec4), - pair_dim=wp.array(mjm.pair_dim, dtype=int), - pair_geom1=wp.array(mjm.pair_geom1, dtype=int), - pair_geom2=wp.array(mjm.pair_geom2, dtype=int), - pair_solref=create_nmodel_batched_array(mjm.pair_solref, dtype=wp.vec2), - pair_solreffriction=create_nmodel_batched_array(mjm.pair_solreffriction, dtype=wp.vec2), - pair_solimp=create_nmodel_batched_array(mjm.pair_solimp, dtype=types.vec5), - pair_margin=create_nmodel_batched_array(mjm.pair_margin, dtype=float), - pair_gap=create_nmodel_batched_array(mjm.pair_gap, dtype=float), - pair_friction=create_nmodel_batched_array(mjm.pair_friction, dtype=types.vec5), - exclude_signature=wp.array(mjm.exclude_signature, dtype=int), - eq_type=wp.array(mjm.eq_type, dtype=int), - eq_obj1id=wp.array(mjm.eq_obj1id, dtype=int), - eq_obj2id=wp.array(mjm.eq_obj2id, dtype=int), - eq_objtype=wp.array(mjm.eq_objtype, dtype=int), - eq_active0=wp.array(mjm.eq_active0, dtype=bool), - eq_solref=create_nmodel_batched_array(mjm.eq_solref, dtype=wp.vec2), - eq_solimp=create_nmodel_batched_array(mjm.eq_solimp, dtype=types.vec5), - eq_data=create_nmodel_batched_array(mjm.eq_data, dtype=types.vec11), - tendon_adr=wp.array(mjm.tendon_adr, dtype=int), - tendon_num=wp.array(mjm.tendon_num, dtype=int), - tendon_limited=wp.array(mjm.tendon_limited, dtype=int), - tendon_actfrclimited=wp.array(mjm.tendon_actfrclimited, dtype=bool), - tendon_solref_lim=create_nmodel_batched_array(mjm.tendon_solref_lim, dtype=wp.vec2f), - tendon_solimp_lim=create_nmodel_batched_array(mjm.tendon_solimp_lim, dtype=types.vec5), - tendon_solref_fri=create_nmodel_batched_array(mjm.tendon_solref_fri, dtype=wp.vec2f), - tendon_solimp_fri=create_nmodel_batched_array(mjm.tendon_solimp_fri, dtype=types.vec5), - tendon_range=create_nmodel_batched_array(mjm.tendon_range, dtype=wp.vec2f), - tendon_actfrcrange=create_nmodel_batched_array(mjm.tendon_actfrcrange, dtype=wp.vec2), - tendon_margin=create_nmodel_batched_array(mjm.tendon_margin, dtype=float), - tendon_stiffness=create_nmodel_batched_array(mjm.tendon_stiffness, dtype=float), - tendon_damping=create_nmodel_batched_array(mjm.tendon_damping, dtype=float), - tendon_armature=create_nmodel_batched_array(mjm.tendon_armature, dtype=float), - tendon_frictionloss=create_nmodel_batched_array(mjm.tendon_frictionloss, dtype=float), - tendon_lengthspring=create_nmodel_batched_array(mjm.tendon_lengthspring, dtype=wp.vec2), - tendon_length0=create_nmodel_batched_array(mjm.tendon_length0, dtype=float), - tendon_invweight0=create_nmodel_batched_array(mjm.tendon_invweight0, dtype=float), - wrap_type=wp.array(mjm.wrap_type, dtype=int), - wrap_objid=wp.array(mjm.wrap_objid, dtype=int), - wrap_prm=wp.array(mjm.wrap_prm, dtype=float), - actuator_trntype=wp.array(mjm.actuator_trntype, dtype=int), - actuator_dyntype=wp.array(mjm.actuator_dyntype, dtype=int), - actuator_gaintype=wp.array(mjm.actuator_gaintype, dtype=int), - actuator_biastype=wp.array(mjm.actuator_biastype, dtype=int), - actuator_trnid=wp.array(mjm.actuator_trnid, dtype=wp.vec2i), - actuator_actadr=wp.array(mjm.actuator_actadr, dtype=int), - actuator_actnum=wp.array(mjm.actuator_actnum, dtype=int), - actuator_ctrllimited=wp.array(mjm.actuator_ctrllimited, dtype=bool), - actuator_forcelimited=wp.array(mjm.actuator_forcelimited, dtype=bool), - actuator_actlimited=wp.array(mjm.actuator_actlimited, dtype=bool), - actuator_dynprm=create_nmodel_batched_array(mjm.actuator_dynprm, dtype=types.vec10f), - actuator_gainprm=create_nmodel_batched_array(mjm.actuator_gainprm, dtype=types.vec10f), - actuator_biasprm=create_nmodel_batched_array(mjm.actuator_biasprm, dtype=types.vec10f), - actuator_actearly=wp.array(mjm.actuator_actearly, dtype=bool), - actuator_ctrlrange=create_nmodel_batched_array(mjm.actuator_ctrlrange, dtype=wp.vec2), - actuator_forcerange=create_nmodel_batched_array(mjm.actuator_forcerange, dtype=wp.vec2), - actuator_actrange=create_nmodel_batched_array(mjm.actuator_actrange, dtype=wp.vec2), - actuator_gear=create_nmodel_batched_array(mjm.actuator_gear, dtype=wp.spatial_vector), - actuator_cranklength=wp.array(mjm.actuator_cranklength, dtype=float), - actuator_acc0=wp.array(mjm.actuator_acc0, dtype=float), - actuator_lengthrange=wp.array(mjm.actuator_lengthrange, dtype=wp.vec2), - sensor_type=wp.array(mjm.sensor_type, dtype=int), - sensor_datatype=wp.array(mjm.sensor_datatype, dtype=int), - sensor_objtype=wp.array(mjm.sensor_objtype, dtype=int), - 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), - plugin=wp.array(plugin_id, dtype=int), - plugin_attr=wp.array(plugin_attr, dtype=wp.vec3f), - M_rownnz=wp.array(mjm.M_rownnz, dtype=int), - M_rowadr=wp.array(mjm.M_rowadr, dtype=int), - M_colind=wp.array(mjm.M_colind, dtype=int), - mapM2M=wp.array(mjm.mapM2M, dtype=int), - # warp only fields: - nacttrnbody=np.sum(mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY), - nsensorcollision=sum(nxn_pairid_collision >= 0), - nsensortaxel=sum(mjm.mesh_vertnum[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_TACTILE]]), - nsensorcontact=np.sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT), - nrangefinder=sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER), - condim_max=condim_max, # TODO(team): get max after filtering, - nmaxpolygon=_compute_nmaxpolygon(mjm, tuple(geom_type_pair_count)), - nmaxmeshdeg=_compute_nmaxmeshdeg(mjm, tuple(geom_type_pair_count)), - has_sdf_geom=bool(np.any(mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF)), - block_dim=types.BlockDim(), - body_tree=body_tree, - mocap_bodyid=wp.array(mocap_bodyid, dtype=int), - body_fluid_ellipsoid=wp.array(body_fluid_ellipsoid, dtype=bool), - # these jnt_limited adrs are used in constraint.py - jnt_limited_slide_hinge_adr=wp.array( - np.nonzero( - mjm.jnt_limited & ((mjm.jnt_type == mujoco.mjtJoint.mjJNT_SLIDE) | (mjm.jnt_type == mujoco.mjtJoint.mjJNT_HINGE)) - )[0], - dtype=int, - ), - jnt_limited_ball_adr=wp.array( - np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0], - dtype=int, - ), - dof_tri_row=wp.array(dof_tri_row, dtype=int), - dof_tri_col=wp.array(dof_tri_col, dtype=int), - geom_pair_type_count=tuple(geom_type_pair_count), - geom_plugin_index=wp.array(geom_plugin_index, dtype=int), - nxn_geom_pair=wp.array(nxn_geom_pair, dtype=wp.vec2i), - nxn_geom_pair_filtered=wp.array(nxn_geom_pair_filtered, dtype=wp.vec2i), - nxn_pairid=wp.array(nxn_pairid, dtype=wp.vec2i), - nxn_pairid_filtered=wp.array(nxn_pairid_filtered, dtype=wp.vec2i), - eq_connect_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.CONNECT)[0], dtype=int), - eq_wld_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.WELD)[0], dtype=int), - eq_jnt_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.JOINT)[0], dtype=int), - eq_ten_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.TENDON)[0], dtype=int), - tendon_jnt_adr=wp.array(tendon_jnt_adr, dtype=int), - tendon_site_pair_adr=wp.array(tendon_site_pair_adr, dtype=int), - tendon_geom_adr=wp.array(tendon_geom_adr, dtype=int), - tendon_limited_adr=wp.array(np.nonzero(mjm.tendon_limited)[0], dtype=int), - ten_wrapadr_site=wp.array(ten_wrapadr_site, dtype=int), - ten_wrapnum_site=wp.array(ten_wrapnum_site, dtype=int), - wrap_jnt_adr=wp.array(wrap_jnt_adr, dtype=int), - wrap_site_adr=wp.array(wrap_site_adr, dtype=int), - wrap_site_pair_adr=wp.array(wrap_site_pair_adr, dtype=int), - wrap_geom_adr=wp.array(wrap_geom_adr, dtype=int), - wrap_pulley_scale=wp.array(wrap_pulley_scale, dtype=float), - actuator_moment_tiles_nv=actuator_moment_tiles_nv, - actuator_moment_tiles_nu=actuator_moment_tiles_nu, - actuator_trntype_body_adr=wp.array(np.nonzero(mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY)[0], dtype=int), - sensor_pos_adr=wp.array( - np.nonzero( - (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS) - & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_JOINTLIMITPOS) - & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONLIMITPOS) - )[0], - dtype=int, - ), - sensor_limitpos_adr=wp.array( - np.nonzero( - (mjm.sensor_type == mujoco.mjtSensor.mjSENS_JOINTLIMITPOS) | (mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONLIMITPOS) - )[0], - dtype=int, - ), - sensor_vel_adr=wp.array( - np.nonzero( - (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL) - & ( - (mjm.sensor_type != mujoco.mjtSensor.mjSENS_JOINTLIMITVEL) - | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONLIMITVEL) - ) - )[0], - dtype=int, - ), - sensor_limitvel_adr=wp.array( - np.nonzero( - (mjm.sensor_type == mujoco.mjtSensor.mjSENS_JOINTLIMITVEL) | (mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONLIMITVEL) - )[0], - dtype=int, - ), - sensor_acc_adr=wp.array( - np.nonzero( - (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC) - & ( - (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TOUCH) - | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_JOINTLIMITFRC) - | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONLIMITFRC) - | (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TENDONACTFRC) - ) - )[0], - dtype=int, - ), - sensor_rangefinder_adr=wp.array(sensor_rangefinder_adr, dtype=int), - rangefinder_sensor_adr=wp.array(rangefinder_sensor_adr, dtype=int), - collision_sensor_adr=wp.array(collision_sensor_adr, dtype=int), - sensor_touch_adr=wp.array( - np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TOUCH)[0], - dtype=int, - ), - sensor_limitfrc_adr=wp.array( - np.nonzero( - (mjm.sensor_type == mujoco.mjtSensor.mjSENS_JOINTLIMITFRC) | (mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONLIMITFRC) - )[0], - dtype=int, - ), - sensor_e_potential=(mjm.sensor_type == mujoco.mjtSensor.mjSENS_E_POTENTIAL).any(), - sensor_e_kinetic=(mjm.sensor_type == mujoco.mjtSensor.mjSENS_E_KINETIC).any(), - sensor_tendonactfrc_adr=wp.array( - np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TENDONACTFRC)[0], - dtype=int, - ), - sensor_subtree_vel=np.isin( - 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_adr_to_contact_adr=wp.array(sensor_adr_to_contact_adr, dtype=int), - sensor_rne_postconstraint=np.isin( - mjm.sensor_type, - [ - mujoco.mjtSensor.mjSENS_ACCELEROMETER, - mujoco.mjtSensor.mjSENS_FORCE, - mujoco.mjtSensor.mjSENS_TORQUE, - mujoco.mjtSensor.mjSENS_FRAMELINACC, - mujoco.mjtSensor.mjSENS_FRAMEANGACC, - ], - ).any(), - sensor_rangefinder_bodyid=wp.array( - mjm.site_bodyid[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER]], dtype=int - ), - sensor_collision_start_adr=wp.array(sensor_collision_start_adr, dtype=int), - 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, - ), - qM_tiles=qM_tiles, - qLD_updates=qLD_updates, - qM_fullm_i=wp.array(qM_fullm_i, dtype=int), - qM_fullm_j=wp.array(qM_fullm_j, dtype=int), - qM_mulm_i=wp.array(qM_mulm_i, dtype=int), - qM_mulm_j=wp.array(qM_mulm_j, dtype=int), - qM_madr_ij=wp.array(qM_madr_ij, dtype=int), - ) + while True: + madr_ij, j = madr_ij + 1, mjm.dof_parentid[j] + if j == -1: + break + m.qM_mulm_i.append(i) + m.qM_mulm_j.append(j) + m.qM_madr_ij.append(madr_ij) + + # place m on device + sizes = dict({"*": 1}, **{f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int}) + for f in dataclasses.fields(types.Model): + if isinstance(f.type, wp.array): + setattr(m, f.name, _create_array(getattr(m, f.name), f.type, sizes)) return m @@ -983,11 +545,7 @@ def _get_padded_sizes(nv: int, njmax: int, is_sparse: bool, tile_size: int): return ((x + multiple - 1) // multiple) * multiple njmax_padded = round_up(njmax, tile_size) - - if is_sparse: - nv_padded = round_up(nv, tile_size) - else: - nv_padded = round_up(nv, 4) + nv_padded = round_up(nv, tile_size) if is_sparse else round_up(nv, 4) return njmax_padded, nv_padded @@ -1031,174 +589,52 @@ def make_data( if njmax < 0: raise ValueError("njmax must be >= 0") - if mujoco.mj_isSparse(mjm): - qM = wp.zeros((nworld, 1, mjm.nM), dtype=float) - qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float) - else: - qM = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float) - qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float) + sizes = dict({"*": 1}, **{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model) if f.type is int}) + sizes["nmaxcondim"] = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max() + sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1)) + tile_size = types.TILE_SIZE_JTDAJ_SPARSE if mujoco.mj_isSparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE + sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, mujoco.mj_isSparse(mjm), tile_size) + sizes["nworld"] = nworld + sizes["naconmax"] = naconmax + sizes["njmax"] = njmax - condim = np.concatenate((mjm.geom_condim, mjm.pair_dim)) - condim_max = np.max(condim) if len(condim) > 0 else 0 + contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)}) + efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)}) + + d_kwargs = { + "contact": contact, + "efc": efc, + "nworld": nworld, + "naconmax": naconmax, + "njmax": njmax, + "qM": None, + "qLD": None, + "geom_xpos": None, + "geom_xmat": None, + } + for f in dataclasses.fields(types.Data): + if f.name in d_kwargs: + continue + d_kwargs[f.name] = _create_array(None, f.type, sizes) + + d = types.Data(**d_kwargs) if mujoco.mj_isSparse(mjm): - tile_size = types.TILE_SIZE_JTDAJ_SPARSE + d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float) + d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float) else: - tile_size = types.TILE_SIZE_JTDAJ_DENSE - - njmax_padded, nv_padded = _get_padded_sizes(mjm.nv, njmax, mujoco.mj_isSparse(mjm), tile_size) + d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float) + d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float) # static geoms (attached to the world) have their poses calculated once during make_data instead # of during each physics step. this speeds up scenes with many static geoms (e.g. terrains) # TODO(team): remove this when we introduce dof islands + sleeping mjd = mujoco.MjData(mjm) mujoco.mj_kinematics(mjm, mjd) - geom_xpos = wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3) - geom_xmat = wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33) + d.geom_xpos = wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3) + d.geom_xmat = wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33) - return types.Data( - solver_niter=wp.zeros(nworld, dtype=int), - ne=wp.zeros(nworld, dtype=int), - nf=wp.zeros(nworld, dtype=int), - nl=wp.zeros(nworld, dtype=int), - nefc=wp.zeros(nworld, dtype=int), - time=wp.zeros(nworld, dtype=float), - energy=wp.zeros(nworld, dtype=wp.vec2), - qpos=wp.zeros((nworld, mjm.nq), dtype=float), - qvel=wp.zeros((nworld, mjm.nv), dtype=float), - act=wp.zeros((nworld, mjm.na), dtype=float), - qacc_warmstart=wp.zeros((nworld, mjm.nv), dtype=float), - ctrl=wp.zeros((nworld, mjm.nu), dtype=float), - qfrc_applied=wp.zeros((nworld, mjm.nv), dtype=float), - xfrc_applied=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - eq_active=wp.array(np.tile(mjm.eq_active0, (nworld, 1)), dtype=bool), - mocap_pos=wp.zeros((nworld, mjm.nmocap), dtype=wp.vec3), - mocap_quat=wp.zeros((nworld, mjm.nmocap), dtype=wp.quat), - qacc=wp.zeros((nworld, mjm.nv), dtype=float), - act_dot=wp.zeros((nworld, mjm.na), dtype=float), - sensordata=wp.zeros((nworld, mjm.nsensordata), dtype=float), - xpos=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), - xquat=wp.zeros((nworld, mjm.nbody), dtype=wp.quat), - xmat=wp.zeros((nworld, mjm.nbody), dtype=wp.mat33), - xipos=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), - ximat=wp.zeros((nworld, mjm.nbody), dtype=wp.mat33), - xanchor=wp.zeros((nworld, mjm.njnt), dtype=wp.vec3), - xaxis=wp.zeros((nworld, mjm.njnt), dtype=wp.vec3), - geom_xpos=geom_xpos, - geom_xmat=geom_xmat, - site_xpos=wp.zeros((nworld, mjm.nsite), dtype=wp.vec3), - site_xmat=wp.zeros((nworld, mjm.nsite), dtype=wp.mat33), - cam_xpos=wp.zeros((nworld, mjm.ncam), dtype=wp.vec3), - cam_xmat=wp.zeros((nworld, mjm.ncam), dtype=wp.mat33), - light_xpos=wp.zeros((nworld, mjm.nlight), dtype=wp.vec3), - light_xdir=wp.zeros((nworld, mjm.nlight), dtype=wp.vec3), - subtree_com=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), - cdof=wp.zeros((nworld, mjm.nv), dtype=wp.spatial_vector), - cinert=wp.zeros((nworld, mjm.nbody), dtype=types.vec10), - flexvert_xpos=wp.zeros((nworld, mjm.nflexvert), dtype=wp.vec3), - flexedge_length=wp.zeros((nworld, mjm.nflexedge), dtype=wp.float32), - ten_wrapadr=wp.zeros((nworld, mjm.ntendon), dtype=int), - ten_wrapnum=wp.zeros((nworld, mjm.ntendon), dtype=int), - ten_J=wp.zeros((nworld, mjm.ntendon, mjm.nv), dtype=float), - ten_length=wp.zeros((nworld, mjm.ntendon), dtype=float), - wrap_obj=wp.zeros((nworld, mjm.nwrap), dtype=wp.vec2i), - wrap_xpos=wp.zeros((nworld, mjm.nwrap), dtype=wp.spatial_vector), - actuator_length=wp.zeros((nworld, mjm.nu), dtype=float), - actuator_moment=wp.zeros((nworld, mjm.nu, mjm.nv), dtype=float), - crb=wp.zeros((nworld, mjm.nbody), dtype=types.vec10), - qM=qM, - qLD=qLD, - qLDiagInv=wp.zeros((nworld, mjm.nv), dtype=float), - flexedge_velocity=wp.zeros((nworld, mjm.nflexedge), dtype=wp.float32), - ten_velocity=wp.zeros((nworld, mjm.ntendon), dtype=float), - actuator_velocity=wp.zeros((nworld, mjm.nu), dtype=float), - cvel=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - cdof_dot=wp.zeros((nworld, mjm.nv), dtype=wp.spatial_vector), - qfrc_bias=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_spring=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_damper=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_gravcomp=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_fluid=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_passive=wp.zeros((nworld, mjm.nv), dtype=float), - subtree_linvel=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), - subtree_angmom=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), - actuator_force=wp.zeros((nworld, mjm.nu), dtype=float), - qfrc_actuator=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_smooth=wp.zeros((nworld, mjm.nv), dtype=float), - qacc_smooth=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_constraint=wp.zeros((nworld, mjm.nv), dtype=float), - qfrc_inverse=wp.zeros((nworld, mjm.nv), dtype=float), - cacc=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - cfrc_int=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - cfrc_ext=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - contact=types.Contact( - dist=wp.zeros((naconmax,), dtype=float), - pos=wp.zeros((naconmax,), dtype=wp.vec3f), - frame=wp.zeros((naconmax,), dtype=wp.mat33f), - includemargin=wp.zeros((naconmax,), dtype=float), - friction=wp.zeros((naconmax,), dtype=types.vec5), - solref=wp.zeros((naconmax,), dtype=wp.vec2f), - solreffriction=wp.zeros((naconmax,), dtype=wp.vec2f), - solimp=wp.zeros((naconmax,), dtype=types.vec5), - dim=wp.zeros((naconmax,), dtype=int), - geom=wp.zeros((naconmax,), dtype=wp.vec2i), - efc_address=wp.zeros((naconmax, np.maximum(1, 2 * (condim_max - 1))), dtype=int), - worldid=wp.zeros((naconmax,), dtype=int), - type=wp.zeros((naconmax,), dtype=int), - geomcollisionid=wp.empty((naconmax,), dtype=int), - ), - efc=types.Constraint( - type=wp.zeros((nworld, njmax), dtype=int), - id=wp.zeros((nworld, njmax), dtype=int), - J=wp.zeros((nworld, njmax_padded, nv_padded), dtype=float), - pos=wp.zeros((nworld, njmax), dtype=float), - margin=wp.zeros((nworld, njmax), dtype=float), - D=wp.zeros((nworld, njmax_padded), dtype=float), - vel=wp.zeros((nworld, njmax), dtype=float), - aref=wp.zeros((nworld, njmax), dtype=float), - frictionloss=wp.zeros((nworld, njmax), dtype=float), - force=wp.zeros((nworld, njmax), dtype=float), - Jaref=wp.zeros((nworld, njmax), dtype=float), - Ma=wp.zeros((nworld, mjm.nv), dtype=float), - grad=wp.zeros((nworld, mjm.nv), dtype=float), - cholesky_L_tmp=wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float), - cholesky_y_tmp=wp.zeros((nworld, mjm.nv), dtype=float), - grad_dot=wp.zeros((nworld,), dtype=float), - Mgrad=wp.zeros((nworld, mjm.nv), dtype=float), - search=wp.zeros((nworld, mjm.nv), dtype=float), - search_dot=wp.zeros((nworld,), dtype=float), - gauss=wp.zeros((nworld,), dtype=float), - cost=wp.zeros((nworld,), dtype=float), - prev_cost=wp.zeros((nworld,), dtype=float), - state=wp.zeros((nworld, njmax_padded), dtype=int), - mv=wp.zeros((nworld, mjm.nv), dtype=float), - jv=wp.zeros((nworld, njmax), dtype=float), - quad=wp.zeros((nworld, njmax), dtype=wp.vec3f), - quad_gauss=wp.zeros((nworld,), dtype=wp.vec3f), - h=wp.zeros((nworld, nv_padded, nv_padded), dtype=float), - alpha=wp.zeros((nworld,), dtype=float), - prev_grad=wp.zeros((nworld, mjm.nv), dtype=float), - prev_Mgrad=wp.zeros((nworld, mjm.nv), dtype=float), - beta=wp.zeros((nworld,), dtype=float), - done=wp.zeros((nworld,), dtype=bool), - ), - # warp only fields: - nworld=nworld, - naconmax=naconmax, - njmax=njmax, - nacon=wp.zeros(1, dtype=int), - ne_connect=wp.zeros(nworld, dtype=int), - ne_weld=wp.zeros(nworld, dtype=int), - ne_jnt=wp.zeros(nworld, dtype=int), - ne_ten=wp.zeros(nworld, dtype=int), - nsolving=wp.zeros(1, dtype=int), - subtree_bodyvel=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - # collision driver - collision_pair=wp.zeros((naconmax,), dtype=wp.vec2i), - collision_pairid=wp.zeros((naconmax,), dtype=wp.vec2i), - collision_worldid=wp.zeros((naconmax,), dtype=int), - ncollision=wp.zeros((1,), dtype=int), - ) + return d def put_data( @@ -1252,265 +688,133 @@ def put_data( if mjd.nefc > njmax: raise ValueError(f"njmax overflow (njmax must be >= {mjd.nefc})") + sizes = dict({"*": 1}, **{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model) if f.type is int}) + sizes["nmaxcondim"] = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max() + sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1)) + tile_size = types.TILE_SIZE_JTDAJ_SPARSE if mujoco.mj_isSparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE + sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, mujoco.mj_isSparse(mjm), tile_size) + sizes["nworld"] = nworld + sizes["naconmax"] = naconmax + sizes["njmax"] = njmax + # ensure static geom positions are computed # TODO: remove once MjData creation semantics are fixed mujoco.mj_kinematics(mjm, mjd) - # calculate some fields that cannot be easily computed inline: + # create contact + contact_kwargs = {"efc_address": None, "worldid": None, "type": None, "geomcollisionid": None} + for f in dataclasses.fields(types.Contact): + if f.name in contact_kwargs: + continue + val = getattr(mjd.contact, f.name) + val = np.repeat(val, nworld, axis=0) + width = ((0, naconmax - val.shape[0]),) + ((0, 0),) * (val.ndim - 1) + val = np.pad(val, width) + contact_kwargs[f.name] = _create_array(val, f.type, sizes) + + contact = types.Contact(**contact_kwargs) + + contact.efc_address = np.zeros((naconmax, sizes["nmaxpyramid"]), dtype=int) + for i in range(mjd.ncon): + efc_address = mjd.contact.efc_address[i] + if efc_address == -1: + continue + condim = mjd.contact.dim[i] + ndim = max(1, 2 * (condim - 1)) if mjm.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL else condim + for j in range(nworld): + contact.efc_address[j * mjd.ncon + i, :ndim] = efc_address + np.arange(ndim) + + contact.efc_address = wp.array(contact.efc_address, dtype=int) + contact.worldid = np.pad(np.repeat(np.arange(nworld), mjd.ncon), (0, naconmax - nworld * mjd.ncon)) + contact.worldid = wp.array(contact.worldid, dtype=int) + contact.type = wp.ones((naconmax,), dtype=int) # TODO(team): set values + contact.geomcollisionid = wp.empty((naconmax,), dtype=int) # TODO(team): set values + + # create efc + efc_kwargs = {"J": None} + + for f in dataclasses.fields(types.Constraint): + if f.name in efc_kwargs: + continue + shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape) + val = np.zeros(shape, dtype=f.type.dtype) + if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"): + val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1)) + efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype) + + efc = types.Constraint(**efc_kwargs) + if mujoco.mj_isSparse(mjm): - qM = np.expand_dims(mjd.qM, axis=0) - qLD = np.expand_dims(mjd.qLD, axis=0) - efc_J = np.zeros((mjd.nefc, mjm.nv)) - mujoco.mju_sparse2dense(efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind) + efc_j = np.zeros((mjd.nefc, mjm.nv)) + mujoco.mju_sparse2dense(efc_j, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind) + else: + efc_j = mjd.efc_J.reshape((mjd.nefc, mjm.nv)) + efc.J = np.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=f.type.dtype) + efc.J[:, : mjd.nefc, : mjm.nv] = np.tile(efc_j, (nworld, 1, 1)) + efc.J = wp.array(efc.J, dtype=float) + + # create data + d_kwargs = { + "contact": contact, + "efc": efc, + "nworld": nworld, + "naconmax": naconmax, + "njmax": njmax, + # fields set after initialization: + "solver_niter": None, + "qM": None, + "qLD": None, + "ten_J": None, + "actuator_moment": None, + "nacon": None, + "ne_connect": None, + "ne_weld": None, + "ne_jnt": None, + "ne_ten": None, + "nsolving": None, + } + for f in dataclasses.fields(types.Data): + if f.name in d_kwargs: + continue + val = getattr(mjd, f.name, None) + if val is not None: + shape = val.shape if hasattr(val, "shape") else () + val = np.full((nworld,) + shape, val) + d_kwargs[f.name] = _create_array(val, f.type, sizes) + + d = types.Data(**d_kwargs) + d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int) + + if mujoco.mj_isSparse(mjm): + d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float) + d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float) ten_J = np.zeros((mjm.ntendon, mjm.nv)) - mujoco.mju_sparse2dense( - ten_J, - mjd.ten_J.reshape(-1), - mjd.ten_J_rownnz, - mjd.ten_J_rowadr, - mjd.ten_J_colind.reshape(-1), - ) + mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1)) + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) else: qM = np.zeros((mjm.nv, mjm.nv)) mujoco.mj_fullM(mjm, qM, mjd.qM) - if (mjd.qM == 0.0).all() or (mjd.qLD == 0.0).all(): - qLD = np.zeros((mjm.nv, mjm.nv)) - else: - qLD = np.linalg.cholesky(qM) - efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv)) + qLD = np.linalg.cholesky(qM) if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv)) + padding = sizes["nv_pad"] - mjm.nv + qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0) + d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float) + d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float) ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv)) + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) # TODO(taylorhowell): sparse actuator_moment actuator_moment = np.zeros((mjm.nu, mjm.nv)) - mujoco.mju_sparse2dense( - actuator_moment, - mjd.actuator_moment, - mjd.moment_rownnz, - mjd.moment_rowadr, - mjd.moment_colind, - ) + mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind) + d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float) - condim = np.concatenate((mjm.geom_condim, mjm.pair_dim)) - condim_max = np.max(condim) if len(condim) > 0 else 0 - contact_efc_address = np.zeros((naconmax, np.maximum(1, 2 * (condim_max - 1))), dtype=int) - for i in range(nworld): - for j in range(mjd.ncon): - condim = mjd.contact.dim[j] - efc_address = mjd.contact.efc_address[j] - if efc_address == -1: - continue - if condim == 1: - nconvar = 1 - else: - nconvar = condim if mjm.opt.cone == mujoco.mjtCone.mjCONE_ELLIPTIC else 2 * (condim - 1) - for k in range(nconvar): - contact_efc_address[i * mjd.ncon + j, k] = efc_address + k + d.nacon = wp.array([mjd.ncon * nworld], dtype=int) + d.ne_connect = wp.full(nworld, 3 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_CONNECT) & mjd.eq_active), dtype=int) + d.ne_weld = wp.full(nworld, 6 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_WELD) & mjd.eq_active), dtype=int) + d.ne_jnt = wp.full(nworld, np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_JOINT) & mjd.eq_active), dtype=int) + d.ne_ten = wp.full(nworld, np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_TENDON) & mjd.eq_active), dtype=int) + d.nsolving = wp.array([nworld], dtype=int) - contact_worldid = np.pad(np.repeat(np.arange(nworld), mjd.ncon), (0, naconmax - nworld * mjd.ncon)) - - ne_connect = int(3 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_CONNECT) & mjd.eq_active)) - ne_weld = int(6 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_WELD) & mjd.eq_active)) - ne_jnt = int(np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_JOINT) & mjd.eq_active)) - ne_ten = int(np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_TENDON) & mjd.eq_active)) - - if mujoco.mj_isSparse(mjm): - tile_size = types.TILE_SIZE_JTDAJ_SPARSE - else: - tile_size = types.TILE_SIZE_JTDAJ_DENSE - - njmax_padded, nv_padded = _get_padded_sizes(mjm.nv, njmax, mujoco.mj_isSparse(mjm), tile_size) - - efc_type_fill = np.zeros((nworld, njmax)) - efc_id_fill = np.zeros((nworld, njmax)) - efc_J_fill = np.zeros((nworld, njmax_padded, nv_padded)) - efc_D_fill = np.zeros((nworld, njmax_padded)) - efc_vel_fill = np.zeros((nworld, njmax)) - efc_pos_fill = np.zeros((nworld, njmax)) - efc_aref_fill = np.zeros((nworld, njmax)) - efc_frictionloss_fill = np.zeros((nworld, njmax)) - efc_force_fill = np.zeros((nworld, njmax)) - efc_margin_fill = np.zeros((nworld, njmax)) - - nefc = mjd.nefc - efc_type_fill[:, :nefc] = np.tile(mjd.efc_type, (nworld, 1)) - efc_id_fill[:, :nefc] = np.tile(mjd.efc_id, (nworld, 1)) - efc_J_fill[:, :nefc, : mjm.nv] = np.tile(efc_J, (nworld, 1, 1)) - efc_D_fill[:, :nefc] = np.tile(mjd.efc_D, (nworld, 1)) - efc_vel_fill[:, :nefc] = np.tile(mjd.efc_vel, (nworld, 1)) - efc_pos_fill[:, :nefc] = np.tile(mjd.efc_pos, (nworld, 1)) - efc_aref_fill[:, :nefc] = np.tile(mjd.efc_aref, (nworld, 1)) - efc_frictionloss_fill[:, :nefc] = np.tile(mjd.efc_frictionloss, (nworld, 1)) - efc_force_fill[:, :nefc] = np.tile(mjd.efc_force, (nworld, 1)) - efc_margin_fill[:, :nefc] = np.tile(mjd.efc_margin, (nworld, 1)) - - # some helper functions to simplify the data field definitions below - - def arr(x, dtype=None): - if not isinstance(x, np.ndarray): - x = np.array(x) - if dtype is None: - if np.issubdtype(x.dtype, np.integer): - dtype = wp.int32 - elif np.issubdtype(x.dtype, np.floating): - dtype = wp.float32 - elif np.issubdtype(x.dtype, bool): - dtype = wp.bool - else: - raise ValueError(f"Unsupported dtype: {x.dtype}") - wp_array = {1: wp.array, 2: wp.array2d, 3: wp.array3d}[x.ndim] - return wp_array(x, dtype=dtype) - - def tile(x, dtype=None): - return arr(np.tile(x, (nworld,) + (1,) * len(x.shape)), dtype) - - def padtile(x, length, dtype=None): - x = np.repeat(x, nworld, axis=0) - width = ((0, length - x.shape[0]),) + ((0, 0),) * (x.ndim - 1) - return arr(np.pad(x, width), dtype) - - return types.Data( - solver_niter=tile(mjd.solver_niter[0]), - ne=wp.full(shape=(nworld), value=mjd.ne), - nf=wp.full(shape=(nworld), value=mjd.nf), - nl=wp.full(shape=(nworld), value=mjd.nl), - nefc=wp.full(shape=(nworld), value=mjd.nefc), - time=arr(mjd.time * np.ones(nworld)), - energy=tile(mjd.energy, dtype=wp.vec2), - qpos=tile(mjd.qpos), - qvel=tile(mjd.qvel), - act=tile(mjd.act), - qacc_warmstart=tile(mjd.qacc_warmstart), - ctrl=tile(mjd.ctrl), - qfrc_applied=tile(mjd.qfrc_applied), - xfrc_applied=tile(mjd.xfrc_applied, dtype=wp.spatial_vector), - eq_active=tile(mjd.eq_active.astype(bool)), - mocap_pos=tile(mjd.mocap_pos, dtype=wp.vec3), - mocap_quat=tile(mjd.mocap_quat, dtype=wp.quat), - qacc=tile(mjd.qacc), - act_dot=tile(mjd.act_dot), - sensordata=tile(mjd.sensordata), - xpos=tile(mjd.xpos, dtype=wp.vec3), - xquat=tile(mjd.xquat, dtype=wp.quat), - xmat=tile(mjd.xmat, dtype=wp.mat33), - xipos=tile(mjd.xipos, dtype=wp.vec3), - ximat=tile(mjd.ximat, dtype=wp.mat33), - xanchor=tile(mjd.xanchor, dtype=wp.vec3), - xaxis=tile(mjd.xaxis, dtype=wp.vec3), - geom_xpos=tile(mjd.geom_xpos, dtype=wp.vec3), - geom_xmat=tile(mjd.geom_xmat, dtype=wp.mat33), - site_xpos=tile(mjd.site_xpos, dtype=wp.vec3), - site_xmat=tile(mjd.site_xmat, dtype=wp.mat33), - cam_xpos=tile(mjd.cam_xpos, dtype=wp.vec3), - cam_xmat=tile(mjd.cam_xmat, dtype=wp.mat33), - light_xpos=tile(mjd.light_xpos, dtype=wp.vec3), - light_xdir=tile(mjd.light_xdir, dtype=wp.vec3), - subtree_com=tile(mjd.subtree_com, dtype=wp.vec3), - cdof=tile(mjd.cdof, dtype=wp.spatial_vector), - cinert=tile(mjd.cinert, dtype=types.vec10), - flexvert_xpos=tile(mjd.flexvert_xpos, dtype=wp.vec3), - flexedge_length=tile(mjd.flexedge_length), - ten_wrapadr=tile(mjd.ten_wrapadr), - ten_wrapnum=tile(mjd.ten_wrapnum), - ten_J=tile(ten_J), - ten_length=tile(mjd.ten_length), - wrap_obj=tile(mjd.wrap_obj, dtype=wp.vec2i), - wrap_xpos=tile(mjd.wrap_xpos, dtype=wp.spatial_vector), - actuator_length=tile(mjd.actuator_length), - actuator_moment=tile(actuator_moment), - crb=tile(mjd.crb, dtype=types.vec10), - qM=tile(qM), - qLD=tile(qLD), - qLDiagInv=tile(mjd.qLDiagInv), - flexedge_velocity=tile(mjd.flexedge_velocity), - ten_velocity=tile(mjd.ten_velocity), - actuator_velocity=tile(mjd.actuator_velocity), - cvel=tile(mjd.cvel, dtype=wp.spatial_vector), - cdof_dot=tile(mjd.cdof_dot, dtype=wp.spatial_vector), - qfrc_bias=tile(mjd.qfrc_bias), - qfrc_spring=tile(mjd.qfrc_spring), - qfrc_damper=tile(mjd.qfrc_damper), - qfrc_gravcomp=tile(mjd.qfrc_gravcomp), - qfrc_fluid=tile(mjd.qfrc_fluid), - qfrc_passive=tile(mjd.qfrc_passive), - subtree_linvel=tile(mjd.subtree_linvel, dtype=wp.vec3), - subtree_angmom=tile(mjd.subtree_angmom, dtype=wp.vec3), - actuator_force=tile(mjd.actuator_force), - qfrc_actuator=tile(mjd.qfrc_actuator), - qfrc_smooth=tile(mjd.qfrc_smooth), - qacc_smooth=tile(mjd.qacc_smooth), - qfrc_constraint=tile(mjd.qfrc_constraint), - qfrc_inverse=tile(mjd.qfrc_inverse), - cacc=tile(mjd.cacc, dtype=wp.spatial_vector), - cfrc_int=tile(mjd.cfrc_int, dtype=wp.spatial_vector), - cfrc_ext=tile(mjd.cfrc_ext, dtype=wp.spatial_vector), - contact=types.Contact( - dist=padtile(mjd.contact.dist, naconmax), - pos=padtile(mjd.contact.pos, naconmax, dtype=wp.vec3), - frame=padtile(mjd.contact.frame, naconmax, dtype=wp.mat33), - includemargin=padtile(mjd.contact.includemargin, naconmax), - friction=padtile(mjd.contact.friction, naconmax, dtype=types.vec5), - solref=padtile(mjd.contact.solref, naconmax, dtype=wp.vec2f), - solreffriction=padtile(mjd.contact.solreffriction, naconmax, dtype=wp.vec2f), - solimp=padtile(mjd.contact.solimp, naconmax, dtype=types.vec5), - dim=padtile(mjd.contact.dim, naconmax), - geom=padtile(mjd.contact.geom, naconmax, dtype=wp.vec2i), - efc_address=arr(contact_efc_address), - worldid=arr(contact_worldid), - type=wp.ones((naconmax,), dtype=int), # TODO(team): set values - geomcollisionid=wp.empty((naconmax,), dtype=int), # TODO(team): set values - ), - efc=types.Constraint( - type=wp.array2d(efc_type_fill, dtype=int), - id=wp.array2d(efc_id_fill, dtype=int), - J=wp.array3d(efc_J_fill, dtype=float), - pos=wp.array2d(efc_pos_fill, dtype=float), - margin=wp.array2d(efc_margin_fill, dtype=float), - D=wp.array2d(efc_D_fill, dtype=float), - vel=wp.array2d(efc_vel_fill, dtype=float), - aref=wp.array2d(efc_aref_fill, dtype=float), - frictionloss=wp.array2d(efc_frictionloss_fill, dtype=float), - force=wp.array2d(efc_force_fill, dtype=float), - Jaref=wp.empty(shape=(nworld, njmax), dtype=float), - Ma=wp.empty(shape=(nworld, mjm.nv), dtype=float), - grad=wp.empty(shape=(nworld, mjm.nv), dtype=float), - cholesky_L_tmp=wp.empty(shape=(nworld, mjm.nv, mjm.nv), dtype=float), - cholesky_y_tmp=wp.empty(shape=(nworld, mjm.nv), dtype=float), - grad_dot=wp.empty(shape=(nworld,), dtype=float), - Mgrad=wp.empty(shape=(nworld, mjm.nv), dtype=float), - search=wp.empty(shape=(nworld, mjm.nv), dtype=float), - search_dot=wp.empty(shape=(nworld,), dtype=float), - gauss=wp.empty(shape=(nworld,), dtype=float), - cost=wp.empty(shape=(nworld,), dtype=float), - prev_cost=wp.empty(shape=(nworld,), dtype=float), - state=wp.zeros(shape=(nworld, njmax_padded), dtype=int), - mv=wp.empty(shape=(nworld, mjm.nv), dtype=float), - jv=wp.empty(shape=(nworld, njmax), dtype=float), - quad=wp.empty(shape=(nworld, njmax), dtype=wp.vec3f), - quad_gauss=wp.empty(shape=(nworld,), dtype=wp.vec3f), - h=wp.zeros(shape=(nworld, nv_padded, nv_padded), dtype=float), - alpha=wp.empty(shape=(nworld,), dtype=float), - 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), - done=wp.empty(shape=(nworld,), dtype=bool), - ), - # warp only fields: - nworld=nworld, - naconmax=naconmax, - njmax=njmax, - nacon=arr([mjd.ncon * nworld]), - ne_connect=wp.full(shape=(nworld), value=ne_connect), - ne_weld=wp.full(shape=(nworld), value=ne_weld), - ne_jnt=wp.full(shape=(nworld), value=ne_jnt), - ne_ten=wp.full(shape=(nworld), value=ne_ten), - nsolving=arr([nworld]), - subtree_bodyvel=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), - # collision driver - collision_pair=wp.empty(naconmax, dtype=wp.vec2i), - collision_pairid=wp.empty(naconmax, dtype=wp.vec2i), - collision_worldid=wp.empty(naconmax, dtype=int), - ncollision=wp.zeros(1, dtype=int), - ) + return d def get_data_into( @@ -1887,7 +1191,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied]) wp.launch( reset_qM, - dim=(d.nworld, 1 if m.opt.is_sparse else m.nv, m.nM if m.opt.is_sparse else m.nv), + dim=(d.nworld, d.qM.shape[1], d.qM.shape[2]), inputs=[reset_input], outputs=[d.qM], ) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py index 459c473b..85a38afb 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -561,7 +561,7 @@ def _flex_elasticity( flex_elem: wp.array(dtype=int), flex_elemedge: wp.array(dtype=int), flexedge_length0: wp.array(dtype=float), - flex_stiffness: wp.array(dtype=float), + flex_stiffness: wp.array2d(dtype=float), flex_damping: wp.array(dtype=float), # Data in: flexvert_xpos_in: wp.array2d(dtype=wp.vec3), @@ -612,8 +612,8 @@ def _flex_elasticity( id = int(0) for ed1 in range(nedge): for ed2 in range(ed1, nedge): - metric[ed1, ed2] = flex_stiffness[21 * elemid + id] - metric[ed2, ed1] = flex_stiffness[21 * elemid + id] + metric[ed1, ed2] = flex_stiffness[elemid, id] + metric[ed2, ed1] = flex_stiffness[elemid, id] id += 1 force = wp.mat(0.0, shape=(6, 3)) @@ -640,7 +640,7 @@ def _flex_bending( flex_vertbodyid: wp.array(dtype=int), flex_edge: wp.array(dtype=wp.vec2i), flex_edgeflap: wp.array(dtype=wp.vec2i), - flex_bending: wp.array(dtype=float), + flex_bending: wp.array2d(dtype=float), # Data in: flexvert_xpos_in: wp.array2d(dtype=wp.vec3), # Data out: @@ -664,7 +664,7 @@ def _flex_bending( return frc = wp.mat(0.0, shape=(4, 3)) - if flex_bending[17 * edgeid + 16]: + if flex_bending[edgeid, 16]: v0 = flexvert_xpos_in[worldid, v[0]] v1 = flexvert_xpos_in[worldid, v[1]] v2 = flexvert_xpos_in[worldid, v[2]] @@ -678,8 +678,8 @@ def _flex_bending( for i in range(nvert): for x in range(3): for j in range(nvert): - force[i, x] -= flex_bending[17 * edgeid + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x] - force[i, x] -= flex_bending[17 * edgeid + 16] * frc[i, x] + force[i, x] -= flex_bending[edgeid, 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x] + force[i, x] -= flex_bending[edgeid, 16] * frc[i, x] for i in range(nvert): bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]] diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py index 0c55b28a..54f1bb48 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -37,12 +37,12 @@ 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.types import vec8f +from mujoco.mjx.third_party.mujoco_warp._src.types import vec8 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) @@ -482,8 +482,8 @@ def _sensor_pos( sensor_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), sensor_pos_adr: wp.array(dtype=int), - sensor_collision_start_adr: wp.array(dtype=int), rangefinder_sensor_adr: wp.array(dtype=int), + sensor_collision_start_adr: wp.array(dtype=int), collision_sensor_adr: wp.array(dtype=int), # Data in: time_in: wp.array(dtype=float), @@ -844,8 +844,8 @@ def sensor_pos(m: Model, d: Data): m.sensor_adr, m.sensor_cutoff, m.sensor_pos_adr, - m.sensor_collision_start_adr, m.rangefinder_sensor_adr, + m.sensor_collision_start_adr, m.collision_sensor_adr, d.time, d.energy, @@ -2089,7 +2089,7 @@ def _sensor_tactile( body_weldid: wp.array(dtype=int), oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8f), + oct_coeff: wp.array(dtype=vec8), geom_type: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), geom_size: wp.array2d(dtype=wp.vec3), 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 5edc76b7..9e56f1f7 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py @@ -35,7 +35,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec10 from mujoco.mjx.third_party.mujoco_warp._src.types import vec11 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) 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 0db3258a..29e7f4ba 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -27,7 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocke from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) @@ -1448,7 +1448,7 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int): @cache_kernel -def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int): +def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int): if njmax < tile_size: tile_size = njmax @@ -1473,7 +1473,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int): nefc = nefc_in[worldid] - sum_val = wp.tile_load(qM_in[worldid], shape=(nv, nv), bounds_check=False) + sum_val = wp.tile_load(qM_in[worldid], shape=(nv_padded, nv_padded), bounds_check=True) # Each tile processes one output tile by looping over all constraints for k in range(0, njmax, TILE_SIZE_K): @@ -1483,7 +1483,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int): # AD: leaving bounds-check disabled here because I'm not entirely sure that # everything always hits the fast path. The padding takes care of any # potential OOB accesses. - J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv), offset=(k, 0), bounds_check=False) + J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv_padded), offset=(k, 0), bounds_check=False) J_kj = J_ki # state check @@ -1499,7 +1499,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int): active_tile = wp.tile_map(active_check, tid_tile, threshold_tile) D_k = wp.tile_map(wp.mul, active_tile, D_k) - J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(nv, TILE_SIZE_K))) + J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(nv_padded, TILE_SIZE_K))) sum_val += wp.tile_matmul(J_ki, J_kj) @@ -1742,8 +1742,9 @@ def _update_gradient(m: types.Model, d: types.Data): outputs=[d.efc.h], ) else: + nv_padded = d.efc.J.shape[2] wp.launch_tiled( - update_gradient_JTDAJ_dense_tiled(m.nv, types.TILE_SIZE_JTDAJ_DENSE, d.njmax), + update_gradient_JTDAJ_dense_tiled(nv_padded, types.TILE_SIZE_JTDAJ_DENSE, d.njmax), dim=d.nworld, inputs=[ d.nefc, 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 a2e3ca8a..f4ec1e94 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -27,7 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) 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 14ceb1b3..1da31efd 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -138,11 +138,13 @@ class DisableBit(enum.IntFlag): DAMPER: passive damper forces GRAVITY: gravitational forces CLAMPCTRL: clamp control to specified range + WARMSTART: warmstart constraint solver + FILTERPARENT: disable collisions between parent and child bodies ACTUATION: apply actuation forces REFSAFE: integrator safety: make ref[0]>=2*timestep + SENSOR: sensors EULERDAMP: implicit damping for Euler integration - FILTERPARENT: disable collisions between parent and child bodies - SENSOR: sensors + NATIVECCD: native convex collision detection (ignored in MJWarp) """ CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT @@ -155,12 +157,12 @@ class DisableBit(enum.IntFlag): GRAVITY = mujoco.mjtDisableBit.mjDSBL_GRAVITY CLAMPCTRL = mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART + FILTERPARENT = mujoco.mjtDisableBit.mjDSBL_FILTERPARENT 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 + EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP + # unsupported: MIDPHASE, AUTORESET, NATIVECCD, ISLAND class EnableBit(enum.IntFlag): @@ -173,7 +175,7 @@ class EnableBit(enum.IntFlag): ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE - # unsupported: OVERRIDE, FWDINV, MULTICCD, ISLAND + # unsupported: OVERRIDE, FWDINV, ISLAND, MULTICCD class TrnType(enum.IntEnum): @@ -594,10 +596,40 @@ class vec11f(wp.types.vector(length=11, dtype=float)): pass +class mat23f(wp.types.matrix(shape=(2, 3), dtype=float)): + pass + + +class mat43f(wp.types.matrix(shape=(4, 3), dtype=float)): + pass + + +class mat63f(wp.types.matrix(shape=(6, 3), dtype=float)): + pass + + vec5 = vec5f vec6 = vec6f +vec8 = vec8f vec10 = vec10f vec11 = vec11f +mat23 = mat23f +mat43 = mat43f +mat63 = mat63f + + +def array(*args) -> wp.array: + """A wrapper around wp.array that adds extra metadata to ease type introspection. + + Format is array(dim_1, dim_2, ..., dtype). dim may be a constant int, or reference a size from + Model or Data (e.g. "nq" or "nworld"). dim may also be "*", which means any nonzero size. + """ + shape, dtype = args[:-1], args[-1] + + arr = wp.array(ndim=len(shape), dtype=dtype) + arr.shape = shape + + return arr @dataclasses.dataclass @@ -641,16 +673,16 @@ class Option: contacts matched after this value is exceded will be ignored """ - timestep: wp.array(dtype=float) - impratio: wp.array(dtype=float) - tolerance: wp.array(dtype=float) - ls_tolerance: wp.array(dtype=float) - ccd_tolerance: wp.array(dtype=float) - density: wp.array(dtype=float) - viscosity: wp.array(dtype=float) - gravity: wp.array(dtype=wp.vec3) - wind: wp.array(dtype=wp.vec3) - magnetic: wp.array(dtype=wp.vec3) + timestep: array("*", float) + impratio: array("*", float) + tolerance: array("*", float) + ls_tolerance: array("*", float) + ccd_tolerance: array("*", float) + density: array("*", float) + viscosity: array("*", float) + gravity: array("*", wp.vec3) + wind: array("*", wp.vec3) + magnetic: array("*", wp.vec3) integrator: int cone: int solver: int @@ -666,8 +698,8 @@ class Option: ls_parallel: bool ls_parallel_min_step: float has_fluid: bool - broadphase: int - broadphase_filter: int + broadphase: BroadphaseType + broadphase_filter: BroadphaseFilter graph_conditional: bool run_collision_detection: bool contact_sensor_maxmatch: int @@ -684,79 +716,6 @@ class Statistic: meaninertia: float -@dataclasses.dataclass -class Constraint: - """Constraint data. - - Attributes: - 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) - margin: inclusion margin (contact) (nworld, njmax) - D: constraint mass (nworld, njmax) - vel: velocity in constraint space: J*qvel (nworld, njmax) - aref: reference pseudo-acceleration (nworld, njmax) - frictionloss: frictionloss (friction) (nworld, njmax) - force: constraint force in constraint space (nworld, njmax) - Jaref: Jac*qacc - aref (nworld, njmax) - Ma: M*qacc (nworld, nv) - grad: gradient of master cost (nworld, nv) - grad_dot: dot(grad, grad) (nworld,) - Mgrad: M / grad (nworld, nv) - search: linesearch vector (nworld, nv) - search_dot: dot(search, search) (nworld,) - gauss: Gauss Cost (nworld,) - cost: constraint + Gauss cost (nworld,) - prev_cost: cost from previous iter (nworld,) - state: constraint state (nworld, njmax) - mv: qM @ search (nworld, nv) - jv: efc_J @ search (nworld, njmax) - quad: quadratic cost coefficients (nworld, njmax, 3) - quad_gauss: quadratic cost Gauss coefficients (nworld, 3) - h: Hessian (nworld, nv, nv) - alpha: line search step size (nworld,) - prev_grad: previous grad (nworld, nv) - prev_Mgrad: previous Mgrad (nworld, nv) - beta: Polak-Ribiere beta (nworld,) - done: solver done (nworld,) - """ - - type: wp.array2d(dtype=int) - id: wp.array2d(dtype=int) - J: wp.array3d(dtype=float) - pos: wp.array2d(dtype=float) - margin: wp.array2d(dtype=float) - D: wp.array2d(dtype=float) - vel: wp.array2d(dtype=float) - aref: wp.array2d(dtype=float) - frictionloss: wp.array2d(dtype=float) - force: wp.array2d(dtype=float) - Jaref: wp.array2d(dtype=float) - Ma: wp.array2d(dtype=float) - grad: wp.array2d(dtype=float) - cholesky_L_tmp: wp.array3d(dtype=float) - cholesky_y_tmp: wp.array2d(dtype=float) - grad_dot: wp.array(dtype=float) - Mgrad: wp.array2d(dtype=float) - search: wp.array2d(dtype=float) - search_dot: wp.array(dtype=float) - gauss: wp.array(dtype=float) - cost: wp.array(dtype=float) - prev_cost: wp.array(dtype=float) - state: wp.array2d(dtype=int) - mv: wp.array2d(dtype=float) - jv: wp.array2d(dtype=float) - quad: wp.array2d(dtype=wp.vec3) - quad_gauss: wp.array(dtype=wp.vec3) - h: wp.array3d(dtype=float) - alpha: wp.array(dtype=float) - prev_grad: wp.array2d(dtype=float) - prev_Mgrad: wp.array2d(dtype=float) - beta: wp.array(dtype=float) - done: wp.array(dtype=bool) - - @dataclasses.dataclass class TileSet: """Tiling configuration for decomposable block diagonal matrix. @@ -782,6 +741,7 @@ class Model: nu: number of actuators/controls na: number of activation states nbody: number of bodies + noct: number of total octree cells in all meshes njnt: number of joints nM: number of non-zeros in sparse inertia matrix nC: number of non-zeros in sparse body-dof matrix @@ -794,7 +754,10 @@ class Model: nflexedge: number of edges in all flexes nflexelem: number of elements in all flexes nflexelemdata: number of element vertex ids in all flexes + nflexelemedge: number of element edge ids in all flexes + nmesh: number of meshes nmeshvert: number of vertices for all meshes + nmeshnormal: number of normals in all meshes nmeshface: number of faces for all meshes nmeshgraph: number of ints in mesh auxiliary data nmeshpoly: number of polygons in all meshes @@ -810,12 +773,13 @@ class Model: nwrap: number of wrap objects in all tendon paths nsensor: number of sensors nmocap: number of mocap bodies + nplugin: number of plugin instances ngravcomp: number of bodies with nonzero gravcomp nsensordata: number of elements in sensor data vector opt: physics options stat: model statistics - qpos0: qpos values at default pose (nworld, nq) - qpos_spring: reference pose for springs (nworld, nq) + qpos0: qpos values at default pose (*, nq) + qpos_spring: reference pose for springs (*, nq) body_parentid: id of body's parent (nbody,) body_rootid: id of root above body (nbody,) body_weldid: id of body that this body is welded to (nbody,) @@ -826,19 +790,19 @@ class Model: body_dofadr: start addr of dofs; -1: no dofs (nbody,) body_geomnum: number of geoms (nbody,) body_geomadr: start addr of geoms; -1: no geoms (nbody,) - body_pos: position offset rel. to parent body (nworld, nbody, 3) - body_quat: orientation offset rel. to parent body (nworld, nbody, 4) - body_ipos: local position of center of mass (nworld, nbody, 3) - body_iquat: local orientation of inertia ellipsoid (nworld, nbody, 4) - body_mass: mass (nworld, nbody,) - body_subtreemass: mass of subtree starting at this body (nworld, nbody,) - body_inertia: diagonal inertia in ipos/iquat frame (nworld, nbody, 3) - body_invweight0: mean inv inert in qpos0 (trn, rot) (nworld, nbody, 2) - body_gravcomp: antigravity force, units of body weight (nworld, nbody) + body_pos: position offset rel. to parent body (*, nbody, 3) + body_quat: orientation offset rel. to parent body (*, nbody, 4) + body_ipos: local position of center of mass (*, nbody, 3) + body_iquat: local orientation of inertia ellipsoid (*, nbody, 4) + body_mass: mass (*, nbody,) + body_subtreemass: mass of subtree starting at this body (*, nbody,) + body_inertia: diagonal inertia in ipos/iquat frame (*, nbody, 3) + body_invweight0: mean inv inert in qpos0 (trn, rot) (*, nbody, 2) + body_gravcomp: antigravity force, units of body weight (*, nbody) body_contype: OR over all geom contypes (nbody,) body_conaffinity: OR over all geom conaffinities (nbody,) oct_child: octree children (noct, 8) - oct_aabb: octree axis-aligned bounding boxes (noct, 6) + oct_aabb: octree axis-aligned bounding boxes (noct, 2, 3) oct_coeff: octree interpolation coefficients (noct, 8) jnt_type: type of joint (JointType) (njnt,) jnt_qposadr: start addr in 'qpos' for joint's data (njnt,) @@ -847,59 +811,59 @@ class Model: jnt_limited: does joint have limits (njnt,) jnt_actfrclimited: does joint have actuator force limits (njnt,) jnt_actgravcomp: is gravcomp force applied via actuators (njnt,) - jnt_solref: constraint solver reference: limit (nworld, njnt, mjNREF) - jnt_solimp: constraint solver impedance: limit (nworld, njnt, mjNIMP) - jnt_pos: local anchor position (nworld, njnt, 3) - jnt_axis: local joint axis (nworld, njnt, 3) - jnt_stiffness: stiffness coefficient (nworld, njnt) - jnt_range: joint limits (nworld, njnt, 2) - jnt_actfrcrange: range of total actuator force (nworld, njnt, 2) - jnt_margin: min distance for limit detection (nworld, njnt) + jnt_solref: constraint solver reference: limit (*, njnt, mjNREF) + jnt_solimp: constraint solver impedance: limit (*, njnt, mjNIMP) + jnt_pos: local anchor position (*, njnt, 3) + jnt_axis: local joint axis (*, njnt, 3) + jnt_stiffness: stiffness coefficient (*, njnt) + jnt_range: joint limits (*, njnt, 2) + jnt_actfrcrange: range of total actuator force (*, njnt, 2) + jnt_margin: min distance for limit detection (*, njnt) dof_bodyid: id of dof's body (nv,) dof_jntid: id of dof's joint (nv,) dof_parentid: id of dof's parent; -1: none (nv,) dof_Madr: dof address in M-diagonal (nv,) - dof_solref: constraint solver reference: frictionloss (nworld, nv, NREF) - dof_solimp: constraint solver impedance: frictionloss (nworld, nv, NIMP) - dof_frictionloss: dof friction loss (nworld, nv) - dof_armature: dof armature inertia/mass (nworld, nv) - dof_damping: damping coefficient (nworld, nv) - dof_invweight0: diag. inverse inertia in qpos0 (nworld, nv) + dof_solref: constraint solver reference: frictionloss (*, nv, NREF) + dof_solimp: constraint solver impedance: frictionloss (*, nv, NIMP) + dof_frictionloss: dof friction loss (*, nv) + dof_armature: dof armature inertia/mass (*, nv) + dof_damping: damping coefficient (*, nv) + dof_invweight0: diag. inverse inertia in qpos0 (*, nv) 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,) geom_bodyid: id of geom's body (ngeom,) geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) - geom_matid: material id for rendering (nworld, ngeom,) + geom_matid: material id for rendering (*, ngeom,) geom_group: geom group inclusion/exclusion mask (ngeom,) geom_priority: geom contact priority (ngeom,) - geom_solmix: mixing coef for solref/imp in geom pair (nworld, ngeom,) - geom_solref: constraint solver reference: contact (nworld, ngeom, mjNREF) - geom_solimp: constraint solver impedance: contact (nworld, ngeom, mjNIMP) - geom_size: geom-specific size parameters (ngeom, 3) - geom_aabb: bounding box, (center, size) (nworld, ngeom, 2, 3) - geom_rbound: radius of bounding sphere (nworld, ngeom,) - geom_pos: local position offset rel. to body (nworld, ngeom, 3) - geom_quat: local orientation offset rel. to body (nworld, ngeom, 4) - geom_friction: friction for (slide, spin, roll) (nworld, ngeom, 3) - geom_margin: detect contact if dist