diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 29cf5d7d..d64f6f3c 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -178,7 +178,7 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any: return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field) # warp arrays - if isinstance(wp_field, wp.types.array): + if isinstance(wp_field, wp.array): return wp_field.numpy() # static @@ -191,7 +191,7 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any: # tuples if isinstance(wp_field, tuple) and len(wp_field) == 0: return () - if isinstance(wp_field, tuple) and isinstance(wp_field[0], wp.types.array): + if isinstance(wp_field, tuple) and isinstance(wp_field[0], wp.array): return tuple(f.numpy() for f in wp_field) if isinstance(wp_field, tuple) and isinstance( wp_field[0], mjwp_types.TileSet diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py index 360803e3..24a01166 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py @@ -15,6 +15,13 @@ """Public API for MJWarp.""" +from importlib import metadata + +try: + __version__ = metadata.version("mujoco_warp") +except metadata.PackageNotFoundError: + __version__ = "unknown" + # isort: off from mujoco.mjx.third_party.mujoco_warp._src.forward import step as step from mujoco.mjx.third_party.mujoco_warp._src.types import Model as Model diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py index f72ceea4..0263928f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py @@ -15,17 +15,12 @@ """Utilities for benchmarking MuJoCo Warp.""" -import importlib -import os import time from typing import Callable, Optional, Tuple -import mujoco import numpy as np import warp as wp -from mujoco.mjx.third_party.mujoco_warp._src import forward -from mujoco.mjx.third_party.mujoco_warp._src import io from mujoco.mjx.third_party.mujoco_warp._src import warp_util from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import Model @@ -165,100 +160,3 @@ def benchmark( run_duration = np.sum(time_vec) return jit_duration, run_duration, trace, nacon, nefc, solver_niter, nsuccess - - -class BenchmarkSuite: - """Base suite for all model benchmarks.""" - - path = "" - batch_size = -1 - nconmax = -1 - njmax = -1 - nstep = 1000 - param_names = ("function",) - params = ( - "jit_duration", - "solver_niter_mean", - "solver_niter_p95", - "device_memory_allocated", - "step", - "step.forward", - "step.forward.fwd_position", - "step.forward.fwd_position.kinematics", - "step.forward.fwd_position.com_pos", - "step.forward.fwd_position.camlight", - "step.forward.fwd_position.crb", - "step.forward.fwd_position.tendon_armature", - "step.forward.fwd_position.collision", - "step.forward.fwd_position.make_constraint", - "step.forward.fwd_position.transmission", - "step.forward.sensor_pos", - "step.forward.fwd_velocity", - "step.forward.fwd_velocity.com_vel", - "step.forward.fwd_velocity.passive", - "step.forward.fwd_velocity.rne", - "step.forward.fwd_velocity.tendon_bias", - "step.forward.sensor_vel", - "step.forward.fwd_actuation", - "step.forward.fwd_acceleration", - "step.forward.fwd_acceleration.xfrc_accumulate", - "step.forward.sensor_acc", - "step.forward.solve", - ) - number = 1 - rounds = 1 - sample_time = 0 - repeat = 1 - replay = "" - - def setup_cache(self): - module = importlib.import_module(self.__module__) - path = os.path.join(os.path.realpath(os.path.dirname(module.__file__)), self.path) - mjm = mujoco.MjModel.from_xml_path(path) - mjd = mujoco.MjData(mjm) - ctrls = None - - if self.replay: - keys = io.find_keys(mjm, self.replay) - if not keys: - raise ValueError(f"Key prefix not find: {self.replay}") - ctrls = io.make_trajectory(mjm, keys) - mujoco.mj_resetDataKeyframe(mjm, mjd, keys[0]) - - if mjm.nkey > 0: - mujoco.mj_resetDataKeyframe(mjm, mjd, 0) - - # TODO(team): mj_forward call shouldn't be necessary, but it is - mujoco.mj_forward(mjm, mjd) - - wp.init() - if os.environ.get("ASV_CACHE_KERNELS", "false").lower() == "false": - wp.clear_kernel_cache() - - free_before = wp.get_device().free_memory - m = io.put_model(mjm) - d = io.put_data(mjm, mjd, self.batch_size, self.nconmax, self.njmax) - free_after = wp.get_device().free_memory - - jit_duration, _, trace, _, _, solver_niter, _ = benchmark(forward.step, m, d, self.nstep, ctrls, True, False, True) - metrics = { - "jit_duration": jit_duration, - "solver_niter_mean": np.mean(solver_niter), - "solver_niter_p95": np.quantile(solver_niter, 0.95), - "device_memory_allocated": free_before - free_after, - } - - def tree_flatten(d, parent_k=""): - ret = {} - steps = self.batch_size * 1000 - for k, v in d.items(): - k = parent_k + "." + k if parent_k else k - ret = ret | {k: 1e6 * v[0][0] / steps} | tree_flatten(v[1], k) - return ret - - metrics = metrics | tree_flatten(trace) - - return metrics - - def track_metric(self, metrics, fn): - return metrics[fn] diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py index 36875caf..e354000a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py @@ -32,9 +32,6 @@ def create_blocked_cholesky_func(block_size: int): It returns a lower-triangular matrix L such that A = L L^T. """ - # TODO(team): remove conditional after mjwarp relies on >= 1.11 - bleeding_edge_warp = wp.static(tuple(map(int, wp.__version__.split(".")[:2])) >= (1, 11)) - # Process the matrix in blocks along its leading dimension. for k in range(0, matrix_size, block_size): end = k + block_size @@ -45,10 +42,7 @@ def create_blocked_cholesky_func(block_size: int): for j in range(0, k, block_size): L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared") - if bleeding_edge_warp: - wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0) - else: - A_kk_tile -= wp.tile_matmul(L_block, wp.tile_transpose(L_block)) + wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0) # Compute the Cholesky factorization for the block L_kk_tile = wp.tile_cholesky(A_kk_tile) @@ -61,15 +55,9 @@ def create_blocked_cholesky_func(block_size: int): for j in range(0, k, block_size): L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared") L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared") - if bleeding_edge_warp: - wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0) - else: - A_ik_tile -= wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile)) + wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0) - if bleeding_edge_warp: - wp.tile_lower_solve_inplace(L_kk_tile, wp.tile_transpose(A_ik_tile)) - else: - A_ik_tile = wp.tile_transpose(wp.tile_lower_solve(L_kk_tile, wp.tile_transpose(A_ik_tile))) + wp.tile_lower_solve_inplace(L_kk_tile, wp.tile_transpose(A_ik_tile)) wp.tile_store(L, A_ik_tile, offset=(i, k)) return blocked_cholesky_func @@ -91,8 +79,6 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int) Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward substitution. """ - # TODO(team): remove conditional after mjwarp relies on >= 1.11 - bleeding_edge_warp = wp.static(tuple(map(int, wp.__version__.split(".")[:2])) >= (1, 11)) rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False) # Forward substitution: solve L y = b @@ -101,17 +87,10 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int) for j in range(0, i, block_size): L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared") y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0)) - if bleeding_edge_warp: - wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0) - else: - rhs_view -= wp.tile_matmul(L_block, y_block) + wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0) L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared") - if bleeding_edge_warp: - wp.tile_lower_solve_inplace(L_tile, rhs_view) - else: - rhs_tmp = wp.tile_lower_solve(L_tile, rhs_view) - wp.tile_assign(rhs_tile, rhs_tmp, offset=(i, 0)) + wp.tile_lower_solve_inplace(L_tile, rhs_view) # Backward substitution: solve L^T x = y for i in range(matrix_size - block_size, -1, -block_size): @@ -120,16 +99,10 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int) for j in range(i_end, matrix_size, block_size): L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared") x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, 0), storage="shared", bounds_check=False) - if bleeding_edge_warp: - wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0) - else: - tmp_tile -= wp.tile_matmul(wp.tile_transpose(L_tile), x_tile) + wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0) L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared") - if bleeding_edge_warp: - wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile) - else: - tmp_tile = wp.tile_upper_solve(wp.tile_transpose(L_tile), tmp_tile) + wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile) wp.tile_store(x, tmp_tile, offset=(i, 0), bounds_check=False) return blocked_cholesky_solve_func 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 e62c6172..e022df89 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 @@ -19,6 +19,7 @@ import warp as wp from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact +from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair @@ -28,8 +29,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR -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 EnableBit 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 @@ -85,8 +86,9 @@ assert _check_convex_collision_pairs(), "_CONVEX_COLLISION_PAIRS is in invalid o @wp.func def _hfield_filter( # Model: + geom_type: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), - geom_aabb: wp.array3d(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), geom_rbound: wp.array2d(dtype=float), geom_margin: wp.array2d(dtype=float), hfield_size: wp.array(dtype=wp.vec4), @@ -109,6 +111,7 @@ def _hfield_filter( # geom info rbound_id = worldid % geom_rbound.shape[0] margin_id = worldid % geom_margin.shape[0] + size_id = worldid % geom_size.shape[0] pos1 = geom_xpos_in[worldid, g1] mat1 = geom_xmat_in[worldid, g1] @@ -135,47 +138,23 @@ def _hfield_filter( mat2 = geom_xmat_in[worldid, g2] mat = mat1T @ mat2 - # aabb for geom in height field frame - xmax = -MJ_MAXVAL - ymax = -MJ_MAXVAL - zmax = -MJ_MAXVAL - xmin = MJ_MAXVAL - ymin = MJ_MAXVAL - zmin = MJ_MAXVAL + # create geom in height field frame for support function queries + geom2 = Geom() + geom2.pos = pos + geom2.rot = mat + geom2.size = geom_size[size_id, g2] + geom2.margin = 0.0 # margin handled separately + geom2.index = -1 - aabb_id = worldid % geom_aabb.shape[0] - center2 = geom_aabb[aabb_id, g2, 0] - size2 = geom_aabb[aabb_id, g2, 1] + geomtype2 = geom_type[g2] - pos += mat1T @ center2 - - sign = wp.vec2(-1.0, 1.0) - - for i in range(2): - for j in range(2): - for k in range(2): - corner_local = wp.vec3(sign[i] * size2[0], sign[j] * size2[1], sign[k] * size2[2]) - corner_hf = mat @ corner_local - - if corner_hf[0] > xmax: - xmax = corner_hf[0] - if corner_hf[1] > ymax: - ymax = corner_hf[1] - if corner_hf[2] > zmax: - zmax = corner_hf[2] - if corner_hf[0] < xmin: - xmin = corner_hf[0] - if corner_hf[1] < ymin: - ymin = corner_hf[1] - if corner_hf[2] < zmin: - zmin = corner_hf[2] - - xmax += pos[0] - xmin += pos[0] - ymax += pos[1] - ymin += pos[1] - zmax += pos[2] - zmin += pos[2] + # use support functions for tight AABB bounds + xmax = support(geom2, geomtype2, wp.vec3(1.0, 0.0, 0.0)).point[0] + xmin = support(geom2, geomtype2, wp.vec3(-1.0, 0.0, 0.0)).point[0] + ymax = support(geom2, geomtype2, wp.vec3(0.0, 1.0, 0.0)).point[1] + ymin = support(geom2, geomtype2, wp.vec3(0.0, -1.0, 0.0)).point[1] + zmax = support(geom2, geomtype2, wp.vec3(0.0, 0.0, 1.0)).point[2] + zmin = support(geom2, geomtype2, wp.vec3(0.0, 0.0, -1.0)).point[2] # box-box test if ( @@ -213,7 +192,6 @@ def ccd_hfield_kernel_builder( geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), - geom_aabb: wp.array3d(dtype=wp.vec3), geom_rbound: wp.array2d(dtype=float), geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array2d(dtype=float), @@ -253,16 +231,13 @@ def ccd_hfield_kernel_builder( collision_worldid_in: wp.array(dtype=int), ncollision_in: wp.array(dtype=int), # In: - epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert1_in: wp.array2d(dtype=wp.vec3), epa_vert2_in: wp.array2d(dtype=wp.vec3), epa_vert_index1_in: wp.array2d(dtype=int), epa_vert_index2_in: wp.array2d(dtype=int), - epa_face_in: wp.array2d(dtype=wp.vec3i), + epa_face_in: wp.array2d(dtype=int), epa_pr_in: wp.array2d(dtype=wp.vec3), epa_norm2_in: wp.array2d(dtype=float), - epa_index_in: wp.array2d(dtype=int), - epa_map_in: wp.array2d(dtype=int), epa_horizon_in: wp.array2d(dtype=int), # Data out: nacon_out: wp.array(dtype=int), @@ -295,7 +270,7 @@ def ccd_hfield_kernel_builder( # height field filter no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter( - geom_dataid, geom_aabb, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2 + geom_type, geom_dataid, geom_size, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2 ) if no_hf_collision: return @@ -346,6 +321,18 @@ def ccd_hfield_kernel_builder( worldid, ) + # transform geom2 into heightfield frame + hf_pos = geom_xpos_in[worldid, g1] + hf_mat = geom_xmat_in[worldid, g1] + hf_matT = wp.transpose(hf_mat) + + geom2.pos = hf_matT @ (geom2.pos - hf_pos) + geom2.rot = hf_matT @ geom2.rot + + # geom1 has identity pose + geom1.pos = wp.vec3(0.0, 0.0, 0.0) + geom1.rot = wp.identity(n=3, dtype=float) + # see MuJoCo mjc_ConvexHField geom1_dataid = geom_dataid[g1] @@ -388,7 +375,6 @@ def ccd_hfield_kernel_builder( 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] @@ -396,8 +382,6 @@ def ccd_hfield_kernel_builder( 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] @@ -405,8 +389,24 @@ def ccd_hfield_kernel_builder( # process all prisms in subgrid count = int(0) for r in range(rmin, rmax): - nvert = int(0) - for c in range(cmin, cmax + 1): + # pre-initialize first 2 vertices + for init_i in range(2): + x = dx * float(cmin) - size[0] + y = dy * float(r + dr[init_i]) - size[1] + z = hfield_data[adr + (r + dr[init_i]) * ncol + cmin] * size[2] + margin + + prism[0] = prism[1] + prism[1] = prism[2] + prism[3] = prism[4] + prism[4] = prism[5] + + prism[2, 0] = x + prism[5, 0] = x + prism[2, 1] = y + prism[5, 1] = y + prism[5, 2] = z + + for c in range(cmin + 1, cmax + 1): # add both triangles from this cell for i in range(2): if count >= MJ_MAXCONPAIR: @@ -432,11 +432,6 @@ def ccd_hfield_kernel_builder( prism[5, 1] = y prism[5, 2] = z - nvert += 1 - - if nvert <= 2: - continue - # prism height test if prism[3, 2] < zmin and prism[4, 2] < zmin and prism[5, 2] < zmin: continue @@ -461,7 +456,6 @@ def ccd_hfield_kernel_builder( geomtype2, x1, geom2.pos, - epa_vert, epa_vert1, epa_vert2, epa_vert_index1, @@ -469,8 +463,6 @@ def ccd_hfield_kernel_builder( epa_face, epa_pr, epa_norm2, - epa_index, - epa_map, epa_horizon, ) @@ -480,13 +472,16 @@ def ccd_hfield_kernel_builder( # cache contact information hfield_contact_dist[count] = dist - pos = 0.5 * (w1 + w2) + # transform contact to global frame + pos_local = 0.5 * (w1 + w2) + pos = hf_mat @ pos_local + hf_pos 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]) + frame_local = make_frame(w1 - w2) + normal_local = wp.vec3(frame_local[0, 0], frame_local[0, 1], frame_local[0, 2]) + normal = hf_mat @ normal_local hfield_contact_normal[count, 0] = normal[0] hfield_contact_normal[count, 1] = normal[1] hfield_contact_normal[count, 2] = normal[2] @@ -720,16 +715,13 @@ def ccd_kernel_builder( # Data in: naconmax_in: int, # In: - epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert1_in: wp.array2d(dtype=wp.vec3), epa_vert2_in: wp.array2d(dtype=wp.vec3), epa_vert_index1_in: wp.array2d(dtype=int), epa_vert_index2_in: wp.array2d(dtype=int), - epa_face_in: wp.array2d(dtype=wp.vec3i), + epa_face_in: wp.array2d(dtype=int), epa_pr_in: wp.array2d(dtype=wp.vec3), epa_norm2_in: wp.array2d(dtype=float), - epa_index_in: wp.array2d(dtype=int), - epa_map_in: wp.array2d(dtype=int), epa_horizon_in: wp.array2d(dtype=int), multiccd_polygon_in: wp.array2d(dtype=wp.vec3), multiccd_clipped_in: wp.array2d(dtype=wp.vec3), @@ -795,7 +787,6 @@ def ccd_kernel_builder( geomtype2, x1, x2, - epa_vert_in[tid], epa_vert1_in[tid], epa_vert2_in[tid], epa_vert_index1_in[tid], @@ -803,8 +794,6 @@ def ccd_kernel_builder( epa_face_in[tid], epa_pr_in[tid], epa_norm2_in[tid], - epa_index_in[tid], - epa_map_in[tid], epa_horizon_in[tid], ) @@ -939,16 +928,13 @@ def ccd_kernel_builder( collision_worldid_in: wp.array(dtype=int), ncollision_in: wp.array(dtype=int), # In: - epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert1_in: wp.array2d(dtype=wp.vec3), epa_vert2_in: wp.array2d(dtype=wp.vec3), epa_vert_index1_in: wp.array2d(dtype=int), epa_vert_index2_in: wp.array2d(dtype=int), - epa_face_in: wp.array2d(dtype=wp.vec3i), + epa_face_in: wp.array2d(dtype=int), epa_pr_in: wp.array2d(dtype=wp.vec3), epa_norm2_in: wp.array2d(dtype=float), - epa_index_in: wp.array2d(dtype=int), - epa_map_in: wp.array2d(dtype=int), epa_horizon_in: wp.array2d(dtype=int), multiccd_polygon_in: wp.array2d(dtype=wp.vec3), multiccd_clipped_in: wp.array2d(dtype=wp.vec3), @@ -1040,7 +1026,6 @@ def ccd_kernel_builder( opt_ccd_tolerance, geom_type, naconmax_in, - epa_vert_in, epa_vert1_in, epa_vert2_in, epa_vert_index1_in, @@ -1048,8 +1033,6 @@ def ccd_kernel_builder( epa_face_in, epa_pr_in, epa_norm2_in, - epa_index_in, - epa_map_in, epa_horizon_in, multiccd_polygon_in, multiccd_clipped_in, @@ -1153,12 +1136,10 @@ def convex_narrowphase(m: Model, d: Data): epa_iterations = m.opt.ccd_iterations # set to true to enable multiccd - use_multiccd = False + use_multiccd = m.opt.enableflags & EnableBit.MULTICCD nmaxpolygon = m.nmaxpolygon if use_multiccd else 0 nmaxmeshdeg = m.nmaxmeshdeg if use_multiccd else 0 - # epa_vert: vertices in EPA polytope in Minkowski space - epa_vert = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3) # epa_vert1: vertices in EPA polytope in geom 1 space epa_vert1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3) # epa_vert2: vertices in EPA polytope in geom 2 space @@ -1168,17 +1149,13 @@ def convex_narrowphase(m: Model, d: Data): # epa_vert_index2: vertex indices in EPA polytope for geom 2 (naconmax, 5 + CCDiter) epa_vert_index2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int) # epa_face: faces of polytope represented by three indices - epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3i) + epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) # epa_pr: projection of origin on polytope faces epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3) # epa_norm2: epa_pr * epa_pr epa_norm2 = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float) - # epa_index: index of face in polytope map - epa_index = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) - # epa_map: status of faces in polytope - epa_map = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int) # epa_horizon: index pair (i j) of edges on horizon - epa_horizon = wp.empty(shape=(d.naconmax, 2 * MJ_MAX_EPAHORIZON), dtype=int) + epa_horizon = wp.empty(shape=(d.naconmax, MJ_MAX_EPAHORIZON), dtype=int) # Contact outputs contact_outputs = [ @@ -1216,7 +1193,6 @@ def convex_narrowphase(m: Model, d: Data): m.geom_solref, m.geom_solimp, m.geom_size, - m.geom_aabb, m.geom_rbound, m.geom_friction, m.geom_margin, @@ -1254,7 +1230,6 @@ def convex_narrowphase(m: Model, d: Data): d.collision_pairid, d.collision_worldid, d.ncollision, - epa_vert, epa_vert1, epa_vert2, epa_vert_index1, @@ -1262,8 +1237,6 @@ def convex_narrowphase(m: Model, d: Data): epa_face, epa_pr, epa_norm2, - epa_index, - epa_map, epa_horizon, ], outputs=contact_outputs, @@ -1342,7 +1315,6 @@ def convex_narrowphase(m: Model, d: Data): d.collision_pairid, d.collision_worldid, d.ncollision, - epa_vert, epa_vert1, epa_vert2, epa_vert_index1, @@ -1350,8 +1322,6 @@ def convex_narrowphase(m: Model, d: Data): epa_face, epa_pr, epa_norm2, - epa_index, - epa_map, epa_horizon, multiccd_polygon, multiccd_clipped, 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 e8ca7e6f..45dd006b 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 @@ -34,6 +34,13 @@ MIN_DIST = 1e-10 FACE_TOL = 0.99999872 EDGE_TOL = 0.00159999931 +# Bit flags for face status in EPA polytope. +# Defined at module scope to avoid Warp's intermediate type issues with literals. +# See: https://github.com/NVIDIA/warp/issues/485 +_FACE_DELETED_BIT = wp.constant(wp.uint32(0x80000000)) +_FACE_INVALID_BIT = wp.constant(wp.uint32(0x40000000)) +_FACE_INVALID_OR_DELETED_MASK = wp.constant(wp.uint32(0xC0000000)) + @wp.struct class GJKResult: @@ -53,7 +60,6 @@ class Polytope: status: int # vertices in polytope - vert: wp.array(dtype=wp.vec3) vert1: wp.array(dtype=wp.vec3) vert2: wp.array(dtype=wp.vec3) vert_index1: wp.array(dtype=int) @@ -61,16 +67,13 @@ class Polytope: nvert: int # faces in polytope - face: wp.array(dtype=wp.vec3i) + # 10 bits per each vertex index, while the last significant bits are for + # invalid and deleted face + face: wp.array(dtype=int) face_pr: wp.array(dtype=wp.vec3) face_norm2: wp.array(dtype=float) - face_index: wp.array(dtype=int) nface: int - # TODO(kbayes): look into if a linear map actually improves performance - face_map: wp.array(dtype=int) - nmap: int - # edges that make up the horizon when adding new vertices to polytope horizon: wp.array(dtype=int) nhorizon: int @@ -91,12 +94,12 @@ def _discrete_geoms(g1: int, g2: int) -> bool: @wp.func -def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: +def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: sp = SupportPoint() sp.cached_index = -1 sp.vertex_index = -1 if geomtype == GeomType.SPHERE: - sp.point = geom.pos + (0.5 * geom.margin) * geom.size[0] * dir + sp.point = geom.pos + (geom.size[0] + 0.5 * geom.margin) * dir return sp local_dir = wp.transpose(geom.rot) @ dir @@ -174,14 +177,13 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: elif geomtype == GeomType.HFIELD: max_dist = float(FLOAT_MIN) # TODO(kbayes): Support edge prisms - sp.vertex_index = wp.where(local_dir[2] < 0, -2, -3) + sp.vertex_index = wp.where(dir[2] < 0, -2, -3) for i in range(6): vert = geom.hfprism[i] - dist = wp.dot(vert, local_dir) + dist = wp.dot(vert, dir) if dist > max_dist: max_dist = dist sp.point = vert - sp.point = geom.rot @ sp.point + geom.pos if geom.margin > 0.0: sp.point += dir * (0.5 * geom.margin) @@ -195,16 +197,15 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float: return 0.0 # compute witness point v - r, ret = _project_origin_plane(pt.vert[v3], pt.vert[v2], pt.vert[v1]) + r, ret = _project_origin_plane(pt.vert1[v3] - pt.vert2[v3], pt.vert1[v2] - pt.vert2[v2], pt.vert1[v1] - pt.vert2[v1]) if ret: return 0.0 - face = wp.vec3i(v1, v2, v3) + face = v1 + (v2 << 10) + (v3 << 20) pt.face[idx] = face pt.face_pr[idx] = r pt.face_norm2[idx] = wp.dot(r, r) - pt.face_index[idx] = -1 return pt.face_norm2[idx] @@ -212,18 +213,16 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float: def _epa_support( pt: Polytope, idx: int, geom1: Geom, geom2: Geom, geom1_type: int, geom2_type: int, dir: wp.vec3 ) -> Tuple[int, int]: - sp = _support(geom1, geom1_type, dir) + sp = support(geom1, geom1_type, dir) pt.vert1[idx] = sp.point pt.vert_index1[idx] = sp.vertex_index index1 = sp.cached_index - sp = _support(geom2, geom2_type, -dir) + sp = support(geom2, geom2_type, -dir) pt.vert2[idx] = sp.point pt.vert_index2[idx] = sp.vertex_index index2 = sp.cached_index - pt.vert[idx] = pt.vert1[idx] - pt.vert2[idx] - return index1, index2 @@ -596,13 +595,13 @@ def gjk( dir_neg = x_k / wp.sqrt(xnorm) # compute kth support point in geom1 - sp = _support(geom1, geomtype1, -dir_neg) + sp = support(geom1, geomtype1, -dir_neg) simplex1[n] = sp.point geom1.index = sp.cached_index simplex_index1[n] = sp.vertex_index # compute kth support point in geom2 - sp = _support(geom2, geomtype2, dir_neg) + sp = support(geom2, geomtype2, dir_neg) simplex2[n] = sp.point geom2.index = sp.cached_index simplex_index2[n] = sp.vertex_index @@ -766,11 +765,6 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult: result = GJKResult() # reset GJK simplex - simplex = mat43() - simplex[0] = pt.vert[v1] - simplex[1] = pt.vert[v2] - simplex[2] = pt.vert[v3] - simplex1 = mat43() simplex1[0] = pt.vert1[v1] simplex1[1] = pt.vert1[v2] @@ -781,6 +775,11 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult: simplex2[1] = pt.vert2[v2] simplex2[2] = pt.vert2[v3] + simplex = mat43() + simplex[0] = simplex1[0] - simplex2[0] + simplex[1] = simplex1[1] - simplex2[1] + simplex[2] = simplex1[2] - simplex2[2] + simplex_index1 = wp.vec4i() simplex_index1[0] = pt.vert_index1[v1] simplex_index1[1] = pt.vert_index1[v2] @@ -835,6 +834,11 @@ def _ray_triangle(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, v4: wp.vec3, v5: wp.vec return 0 +@wp.func +def _get_edge(edge: int) -> wp.vec2i: + return wp.vec2i(edge & 0x3FF, (edge >> 10) & 0x3FF) + + @wp.func def _add_edge(pt: Polytope, e1: int, e2: int) -> int: n = pt.nhorizon @@ -842,47 +846,30 @@ def _add_edge(pt: Polytope, e1: int, e2: int) -> int: if n < 0: return -1 + edge = (wp.min(e1, e2) << 10) | wp.max(e1, e2) + for i in range(n): - old_e1 = pt.horizon[2 * i + 0] - old_e2 = pt.horizon[2 * i + 1] - if (old_e1 == e1 and old_e2 == e2) or (old_e1 == e2 and old_e2 == e1): - pt.horizon[2 * i + 0] = pt.horizon[2 * (n - 1) + 0] - pt.horizon[2 * i + 1] = pt.horizon[2 * (n - 1) + 1] + if edge == pt.horizon[i]: + pt.horizon[i] = pt.horizon[n - 1] return n - 1 # out of memory, force EPA to return early without contact - if n > pt.horizon.shape[0] - 2: + if n == pt.horizon.shape[0]: return -1 - pt.horizon[2 * n + 0] = e1 - pt.horizon[2 * n + 1] = e2 + pt.horizon[n] = edge return n + 1 -@wp.func -def _delete_face(pt: Polytope, face_id: int) -> int: - index = pt.face_index[face_id] - # delete from map - if index >= 0: - last_face = pt.face_map[pt.nmap - 1] - pt.face_map[index] = last_face - pt.face_index[last_face] = index - pt.nmap -= 1 - # mark face as deleted from polytope - pt.face_index[face_id] = -2 - return pt.nmap - - @wp.func def _epa_witness( pt: Polytope, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int, face_idx: int ) -> Tuple[wp.vec3, wp.vec3, float]: - face = pt.face[face_idx] - + face = _get_face_verts(pt.face[face_idx]) # compute affine coordinates for witness points on plane defined by face - v1 = pt.vert[face[0]] - v2 = pt.vert[face[1]] - v3 = pt.vert[face[2]] + v1 = pt.vert1[face[0]] - pt.vert2[face[0]] + v2 = pt.vert1[face[1]] - pt.vert2[face[1]] + v3 = pt.vert1[face[2]] - pt.vert2[face[2]] coordinates = _tri_affine_coord(v1, v2, v3, pt.face_pr[face_idx]) l1 = coordinates[0] @@ -904,12 +891,12 @@ def _epa_witness( i3 = pt.vert_index1[face[2]] if geomtype1 == GeomType.HFIELD and (i1 != i2 or i1 != i3): # TODO(kbayes): Fix case where geom2 is near bottom of height field or "extreme" prism heights - n = geom1.rot[:, 2] + n = wp.vec3(0.0, 0.0, 1.0) - # height field prism vertices in global frame - a = geom1.pos + geom1.rot @ geom1.hfprism[3] - b = geom1.pos + geom1.rot @ geom1.hfprism[4] - c = geom1.pos + geom1.rot @ geom1.hfprism[5] + # height field prism vertices + a = geom1.hfprism[3] + b = geom1.hfprism[4] + c = geom1.hfprism[5] # TODO(kbayes): Support cases where geom2 is larger than the height field if geomtype2 == GeomType.CAPSULE or geomtype2 == GeomType.SPHERE: @@ -917,13 +904,13 @@ def _epa_witness( margin = geom2.margin geom2.margin = 0.0 geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2]) - sp = _support(geom2, geomtype2, x2) + sp = support(geom2, geomtype2, x2) x2 = sp.point - (0.5 * margin + radius) * n geom2.size[0] = radius geom2.margin = margin else: x2 = wp.normalize(x2) - sp = _support(geom2, geomtype2, x2) + sp = support(geom2, geomtype2, x2) x2 = sp.point coordinates2 = _tri_affine_coord(a, b, c, x2) @@ -952,7 +939,6 @@ def _epa_witness( def _polytope2( # In: pt: Polytope, - dist: float, simplex: mat43, simplex1: mat43, simplex2: mat43, @@ -985,9 +971,6 @@ def _polytope2( d3 = R @ d2 # save vertices and get indices for each one - pt.vert[0] = simplex[0] - pt.vert[1] = simplex[1] - pt.vert1[0] = simplex1[0] pt.vert1[1] = simplex1[1] @@ -1030,19 +1013,16 @@ def _polytope2( return pt, _replace_simplex3(pt, 1, 4, 3) # check hexahedron is convex - if not _ray_triangle(simplex[0], simplex[1], pt.vert[2], pt.vert[3], pt.vert[4]): + v2 = pt.vert1[2] - pt.vert2[2] + v3 = pt.vert1[3] - pt.vert2[3] + v4 = pt.vert1[4] - pt.vert2[4] + if not _ray_triangle(simplex[0], simplex[1], v2, v3, v4): pt.status = 1 return pt, GJKResult() - # populate face map - for i in range(6): - pt.face_map[i] = i - pt.face_index[i] = i - # set polytope counts pt.nvert = 5 pt.nface = 6 - pt.nmap = 6 pt.status = 0 return pt, GJKResult() @@ -1069,10 +1049,6 @@ def _polytope3( pt.status = 2 return pt - pt.vert[0] = simplex[0] - pt.vert[1] = simplex[1] - pt.vert[2] = simplex[2] - pt.vert1[0] = simplex1[0] pt.vert1[1] = simplex1[1] pt.vert1[2] = simplex1[2] @@ -1095,8 +1071,8 @@ def _polytope3( v1 = simplex[0] v2 = simplex[1] v3 = simplex[2] - v4 = pt.vert[3] - v5 = pt.vert[4] + v4 = pt.vert1[3] - pt.vert2[3] + v5 = pt.vert1[4] - pt.vert2[4] # check that v4 is not contained in the 2-simplex if _tri_point_intersect(v1, v2, v3, v4): @@ -1134,15 +1110,9 @@ def _polytope3( pt.status = 11 return pt - # populate face map - for i in range(6): - pt.face_map[i] = i - pt.face_index[i] = i - # set polytope counts pt.nvert = 5 pt.nface = 6 - pt.nmap = 6 pt.status = 0 return pt @@ -1151,23 +1121,13 @@ def _polytope3( def _polytope4( # In: pt: Polytope, - dist: float, simplex: mat43, simplex1: mat43, simplex2: mat43, simplex_index1: wp.vec4i, simplex_index2: wp.vec4i, - geom1: Geom, - geom2: Geom, - geomtype1: int, - geomtype2: int, ) -> Tuple[Polytope, GJKResult]: """Create polytope for EPA given a 3-simplex from GJK.""" - pt.vert[0] = simplex[0] - pt.vert[1] = simplex[1] - pt.vert[2] = simplex[2] - pt.vert[3] = simplex[3] - pt.vert1[0] = simplex1[0] pt.vert1[1] = simplex1[1] pt.vert1[2] = simplex1[2] @@ -1205,23 +1165,47 @@ def _polytope4( pt.status = -1 return pt, _replace_simplex3(pt, 3, 2, 1) - if not _test_tetra(pt.vert[0], pt.vert[1], pt.vert[2], pt.vert[3]): + if not _test_tetra(simplex[0], simplex[1], simplex[2], simplex[3]): pt.status = 12 return pt, GJKResult() - # populate face map - for i in range(4): - pt.face_map[i] = i - pt.face_index[i] = i - # set polytope counts pt.nvert = 4 pt.nface = 4 - pt.nmap = 4 pt.status = 0 return pt, GJKResult() +@wp.func +def _get_face_verts(face: int) -> wp.vec3i: + """Return the three vertices of the face given by indices into the polytope vertex array.""" + return wp.vec3i(face & 0x3FF, face >> 10 & 0x3FF, face >> 20 & 0x3FF) + + +@wp.func +def _delete_face(face: int) -> int: + """Return the face with the deleted bit enabled.""" + return int(wp.uint32(face) | _FACE_DELETED_BIT) + + +@wp.func +def _is_face_deleted(face: int) -> bool: + """Return true if face is deleted.""" + return bool(wp.uint32(face) & _FACE_DELETED_BIT) + + +@wp.func +def _invalidate_face(face: int) -> int: + """Return the face with the invalid bit enabled.""" + return int(wp.uint32(face) | _FACE_INVALID_BIT) + + +@wp.func +def _is_invalid_face(face: int) -> bool: + """Return true if face is invalid or deleted.""" + return bool(wp.uint32(face) & _FACE_INVALID_OR_DELETED_MASK) + + @wp.func def _epa( # In: @@ -1242,18 +1226,22 @@ def _epa( pidx = int(-1) epsilon = wp.where(is_discrete, 1e-15, tolerance) cnt = int(1) + nvalid = pt.nface # number of potential faces for expanding the polytope + # the face vertices are encoded in 10-bits that index the vertex array, + # so iterations must be cap to limit the number of generated vertices + # (one new vertex per iteration) + epa_iterations = wp.min(epa_iterations, 1000) for _ in range(epa_iterations): pidx = idx idx = int(-1) + lower2 = float(FLOAT_MAX) # find the face closest to the origin (lower bound for penetration depth) - lower2 = float(FLOAT_MAX) - for i in range(pt.nmap): - face_idx = pt.face_map[i] - if pt.face_norm2[face_idx] < lower2: - idx = int(face_idx) - lower2 = float(pt.face_norm2[face_idx]) + for i in range(pt.nface): + if not _is_invalid_face(pt.face[i]) and pt.face_norm2[i] < lower2: + idx = i + lower2 = pt.face_norm2[i] # face not valid, return previous face if lower2 > upper2 or idx < 0: @@ -1269,12 +1257,13 @@ def _epa( wi = pt.nvert face_pr_normalized = pt.face_pr[idx] / lower i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr_normalized) + w = pt.vert1[wi] - pt.vert2[wi] geom1.index = i1 geom2.index = i2 pt.nvert += 1 # upper bound for kth iteration - upper_k = wp.dot(face_pr_normalized, pt.vert[wi]) + upper_k = wp.dot(face_pr_normalized, w) if upper_k < upper: upper = upper_k upper2 = upper * upper @@ -1292,45 +1281,51 @@ def _epa( if found_repeated: break - pt.nmap = _delete_face(pt, idx) - pt.nhorizon = _add_edge(pt, pt.face[idx][0], pt.face[idx][1]) - pt.nhorizon = _add_edge(pt, pt.face[idx][1], pt.face[idx][2]) - pt.nhorizon = _add_edge(pt, pt.face[idx][2], pt.face[idx][0]) + nvalid -= 1 + pt.face[idx] = _delete_face(pt.face[idx]) + face = _get_face_verts(pt.face[idx]) + pt.nhorizon = _add_edge(pt, face[0], face[1]) + pt.nhorizon = _add_edge(pt, face[1], face[2]) + pt.nhorizon = _add_edge(pt, face[2], face[0]) if pt.nhorizon == -1: + wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0]) idx = -1 break # compute horizon for w for i in range(pt.nface): - if pt.face_index[i] == -2: + if _is_face_deleted(pt.face[i]): continue - if wp.dot(pt.face_pr[i], pt.vert[wi]) - pt.face_norm2[i] > 1e-10: - pt.nmap = _delete_face(pt, i) - pt.nhorizon = _add_edge(pt, pt.face[i][0], pt.face[i][1]) - pt.nhorizon = _add_edge(pt, pt.face[i][1], pt.face[i][2]) - pt.nhorizon = _add_edge(pt, pt.face[i][2], pt.face[i][0]) + if wp.dot(pt.face_pr[i], w) - pt.face_norm2[i] > 1e-10: + nvalid = wp.where(_is_invalid_face(pt.face[i]), nvalid, nvalid - 1) + pt.face[i] = _delete_face(pt.face[i]) + face = _get_face_verts(pt.face[i]) + pt.nhorizon = _add_edge(pt, face[0], face[1]) + pt.nhorizon = _add_edge(pt, face[1], face[2]) + pt.nhorizon = _add_edge(pt, face[2], face[0]) if pt.nhorizon == -1: + wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0]) idx = -1 break # insert w as new vertex and attach faces along the horizon for i in range(pt.nhorizon): - dist2 = _attach_face(pt, pt.nface, wi, pt.horizon[2 * i + 0], pt.horizon[2 * i + 1]) + edge = _get_edge(pt.horizon[i]) + dist2 = _attach_face(pt, pt.nface, wi, edge[0], edge[1]) if dist2 == 0: idx = -1 break pt.nface += 1 - # store face in map if dist2 >= lower2 and dist2 <= upper2: - pt.face_map[pt.nmap] = pt.nface - 1 - pt.face_index[pt.nface - 1] = pt.nmap - pt.nmap += 1 + nvalid += 1 + else: + pt.face[pt.nface - 1] = _invalidate_face(pt.face[pt.nface - 1]) # no face candidates left - if pt.nmap == 0 or idx == -1: + if nvalid == 0 or idx == -1: break # clear horizon @@ -1394,7 +1389,7 @@ def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i: if c == b: c = _next(npolygon, c) if d == c: - d == _next(npolygon, d) + d = _next(npolygon, d) return res @@ -1889,7 +1884,7 @@ def _polygon_clip( for i in range(npolygon): # get edge PQ of the polygon P = polygon_out[i] - Q = wp.where(i < npolygon - 1, polygon_out[i + 1], polygon_out[0]) + Q = polygon_out[(i + 1) % npolygon] # determine if P and Q are in the halfspace of the clipping edge inside1 = _halfspace(face1[e], pn[e], P) @@ -1968,7 +1963,7 @@ def multicontact( epa_vert2: wp.array(dtype=wp.vec3), epa_vert_index1: wp.array(dtype=int), epa_vert_index2: wp.array(dtype=int), - face: wp.vec3i, + epa_face: int, x1: wp.vec3, x2: wp.vec3, geom1: Geom, @@ -2000,6 +1995,8 @@ def multicontact( polymapnum = geom2.mesh_polymapnum polymap = geom2.mesh_polymap + face = _get_face_verts(epa_face) + # get dimensions of features of geoms 1 and 2 nface1, feature_index1, feature_vertex1 = _feature_dim(face, epa_vert_index1, epa_vert1) nface2, feature_index2, feature_vertex2 = _feature_dim(face, epa_vert_index2, epa_vert2) @@ -2190,13 +2187,15 @@ def _inflate( break if is_side: - n = geom1.rot[:, 2] - sp = _support(geom2, geomtype2, x2) + n = wp.vec3(0.0, 0.0, 1.0) + sp = support(geom2, geomtype2, x2) x2 = sp.point - margin2 * n + # height field prism vertices a = geom1.hfprism[3] b = geom1.hfprism[4] c = geom1.hfprism[5] + coordinates = _tri_affine_coord(a, b, c, x2) if coordinates[0] > 0 and coordinates[1] > 0 and coordinates[2] > 0: x1 = coordinates[0] * a + coordinates[1] * b + coordinates[2] * c @@ -2231,16 +2230,13 @@ def ccd( geomtype2: int, x_1: wp.vec3, x_2: wp.vec3, - vert: wp.array(dtype=wp.vec3), vert1: wp.array(dtype=wp.vec3), vert2: wp.array(dtype=wp.vec3), vert_index1: wp.array(dtype=int), vert_index2: wp.array(dtype=int), - face: wp.array(dtype=wp.vec3i), + face: wp.array(dtype=int), face_pr: wp.array(dtype=wp.vec3), face_norm2: wp.array(dtype=float), - face_index: wp.array(dtype=int), - face_map: wp.array(dtype=int), horizon: wp.array(dtype=int), ) -> Tuple[float, int, wp.vec3, wp.vec3, int]: """General convex collision detection via GJK/EPA.""" @@ -2291,10 +2287,8 @@ def ccd( pt = Polytope() pt.nface = 0 - pt.nmap = 0 pt.nvert = 0 pt.nhorizon = 0 - pt.vert = vert pt.vert1 = vert1 pt.vert2 = vert2 pt.vert_index1 = vert_index1 @@ -2302,14 +2296,11 @@ def ccd( pt.face = face pt.face_pr = face_pr pt.face_norm2 = face_norm2 - pt.face_index = face_index - pt.face_map = face_map pt.horizon = horizon if result.dim == 2: pt, new_result = _polytope2( pt, - result.dist, result.simplex, result.simplex1, result.simplex2, @@ -2330,16 +2321,11 @@ def ccd( elif result.dim == 4: pt, new_result = _polytope4( pt, - result.dist, result.simplex, result.simplex1, result.simplex2, result.simplex_index1, result.simplex_index2, - geom1, - geom2, - geomtype1, - geomtype2, ) if pt.status == -1: result.simplex = new_result.simplex 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 5cf03c7d..ff6d80a1 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 @@ -822,39 +822,41 @@ def capsule_capsule_wrapper( cap2_axis, cap2.size[0], # radius2 cap2.size[1], # half_length2 + margin, ) - write_contact( - naconmax_in, - 0, - dist, - pos, - make_frame(normal), - margin, - gap, - condim, - friction, - solref, - solreffriction, - solimp, - geoms, - 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, - ) + for i in range(2): + write_contact( + naconmax_in, + i, + dist[i], + wp.vec3(pos[i, 0], pos[i, 1], pos[i, 2]), + make_frame(wp.vec3(normal[i, 0], normal[i, 1], normal[i, 2])), + margin, + gap, + condim, + friction, + solref, + solreffriction, + solimp, + geoms, + 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, + ) @wp.func diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py index 04d2ea3c..a54565da 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py @@ -188,7 +188,8 @@ def capsule_capsule( cap2_axis: wp.vec3, cap2_radius: float, cap2_half_length: float, -) -> Tuple[float, wp.vec3, wp.vec3]: + margin: float, +) -> Tuple[wp.vec2, mat23f, mat23f]: """Core contact geometry calculation for capsule-capsule collision. Args: @@ -200,28 +201,110 @@ def capsule_capsule( cap2_axis: Axis direction of the second capsule. cap2_radius: Radius of the second capsule. cap2_half_length: Half length of the second capsule. + margin: Collision margin for filtering contacts. Returns: - - Vector of contact distances. + - Vector of contact distances (wp.inf for invalid contacts). - Matrix of contact positions (one per row). - Matrix of contact normal vectors (one per row). """ - # TODO(team): parallel axes case + contact_dist = wp.vec2(wp.inf, wp.inf) + contact_pos = mat23f() + contact_normal = mat23f() - # Calculate capsule segments - seg1 = cap1_axis * cap1_half_length - seg2 = cap2_axis * cap2_half_length + # calculate scaled axes and center difference + axis1 = cap1_axis * cap1_half_length + axis2 = cap2_axis * cap2_half_length + dif = cap1_pos - cap2_pos - # Find closest points between capsule centerlines - pt1, pt2 = closest_segment_to_segment_points( - cap1_pos - seg1, - cap1_pos + seg1, - cap2_pos - seg2, - cap2_pos + seg2, - ) + # compute matrix coefficients and determinant + ma = wp.dot(axis1, axis1) + mb = -wp.dot(axis1, axis2) + mc = wp.dot(axis2, axis2) + u = -wp.dot(axis1, dif) + v = wp.dot(axis2, dif) + det = ma * mc - mb * mb - # Use sphere-sphere collision between closest points - return sphere_sphere(pt1, cap1_radius, pt2, cap2_radius) + # non-parallel axes: 1 contact + if wp.abs(det) >= MJ_MINVAL: + inv_det = 1.0 / det + x1 = (mc * u - mb * v) * inv_det + x2 = (ma * v - mb * u) * inv_det + + if x1 > 1.0: + x1 = 1.0 + x2 = (v - mb) / mc + elif x1 < -1.0: + x1 = -1.0 + x2 = (v + mb) / mc + + if x2 > 1.0: + x2 = 1.0 + x1 = wp.clamp((u - mb) / ma, -1.0, 1.0) + elif x2 < -1.0: + x2 = -1.0 + x1 = wp.clamp((u + mb) / ma, -1.0, 1.0) + + # find nearest points + vec1 = cap1_pos + axis1 * x1 + vec2 = cap2_pos + axis2 * x2 + + dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius) + if dist <= margin: + contact_dist[0] = dist + contact_pos[0] = pos + contact_normal[0] = normal + + # parallel axes: test all 4 endpoint pairs, keep first 2 that pass margin check + else: + contact_count = 0 + + # x1 = 1: test positive end of capsule 1 + vec1 = cap1_pos + axis1 + x2 = wp.clamp((v - mb) / mc, -1.0, 1.0) + vec2 = cap2_pos + axis2 * x2 + dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius) + if dist <= margin: + contact_dist[contact_count] = dist + contact_pos[contact_count] = pos + contact_normal[contact_count] = normal + contact_count += 1 + + # x1 = -1: test negative end of capsule 1 + vec1 = cap1_pos - axis1 + x2 = wp.clamp((v + mb) / mc, -1.0, 1.0) + vec2 = cap2_pos + axis2 * x2 + dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius) + if dist <= margin: + contact_dist[contact_count] = dist + contact_pos[contact_count] = pos + contact_normal[contact_count] = normal + contact_count += 1 + + # x2 = 1: test positive end of capsule 2 + if contact_count < 2: + vec2 = cap2_pos + axis2 + x1 = wp.clamp((u - mb) / ma, -1.0, 1.0) + vec1 = cap1_pos + axis1 * x1 + dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius) + if dist <= margin: + contact_dist[contact_count] = dist + contact_pos[contact_count] = pos + contact_normal[contact_count] = normal + contact_count += 1 + + # x2 = -1: test negative end of capsule 2 + if contact_count < 2: + vec2 = cap2_pos - axis2 + x1 = wp.clamp((u + mb) / ma, -1.0, 1.0) + vec1 = cap1_pos + axis1 * x1 + dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius) + if dist <= margin: + contact_dist[contact_count] = dist + contact_pos[contact_count] = pos + contact_normal[contact_count] = normal + + return contact_dist, contact_pos, contact_normal @wp.func 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 f9392cad..65acf580 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 @@ -66,7 +66,6 @@ class MeshData: mesh_faceadr: wp.array(dtype=int) mesh_face: wp.array(dtype=wp.vec3i) data_id: int - data_id: int pos: wp.vec3 mat: wp.mat33 pnt: wp.vec3 @@ -169,7 +168,7 @@ def grad_sphere(p: wp.vec3) -> wp.vec3: if c > 1e-9: return p / c else: - wp.vec3(0.0) + return wp.vec3(0.0) @wp.func 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 4f8a0501..82f587da 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -486,8 +486,8 @@ def implicit(m: Model, d: Data): qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float) qLD = wp.empty((d.nworld, 1, m.nC), dtype=float) else: - qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float) - qLD = wp.empty((d.nworld, m.nv, m.nv), dtype=float) + qDeriv = wp.empty(d.qM.shape, dtype=float) + qLD = wp.empty(d.qM.shape, dtype=float) qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float) derivative.deriv_smooth_vel(m, d, qDeriv) qacc = wp.empty((d.nworld, m.nv), dtype=float) @@ -956,6 +956,7 @@ def step1(m: Model, d: Data): # TODO(team): mj_checkPos # TODO(team): mj_checkVel fwd_position(m, d) + d.sensordata.zero_() sensor.sensor_pos(m, d) if energy: 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 c7e00d2a..60ac1647 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -14,6 +14,9 @@ # ============================================================================== import dataclasses +import importlib.metadata +import re +import warnings from typing import Any, Optional, Sequence, Union import mujoco @@ -25,6 +28,24 @@ from mujoco.mjx.third_party.mujoco_warp._src import warp_util from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel +def _is_mujoco_dev() -> bool: + _DEV_VERSION_PATTERN = re.compile(r"^\d+\.\d+\.\d+.+") # anything after x.y.z + + version = getattr(__import__("mujoco"), "__version__", None) + if version and _DEV_VERSION_PATTERN.match(version): + return True + + # fall back to metadata + dist_version = importlib.metadata.version("mujoco") + if _DEV_VERSION_PATTERN.match(dist_version): + return True + + return False + + +BLEEDING_EDGE_MUJOCO = _is_mujoco_dev() + + 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. @@ -72,36 +93,38 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: warp_util.check_toolkit_driver() # 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), + for field, field_type, mj_type in ( + (mjm.actuator_trntype, types.TrnType, mujoco.mjtTrn), + (mjm.actuator_dyntype, types.DynType, mujoco.mjtDyn), + (mjm.actuator_gaintype, types.GainType, mujoco.mjtGain), + (mjm.actuator_biastype, types.BiasType, mujoco.mjtBias), + (mjm.eq_type, types.EqType, mujoco.mjtEq), + (mjm.geom_type, types.GeomType, mujoco.mjtGeom), + (mjm.sensor_type, types.SensorType, mujoco.mjtSensor), + (mjm.wrap_type, types.WrapType, mujoco.mjtWrap), ): missing = ~np.isin(field, field_type) if missing.any(): - raise NotImplementedError(f"{field_type.__name__}: {field[missing]} not supported.") + names = [mj_type(v).name for v in field[missing]] + raise NotImplementedError(f"{names} not supported.") # 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), + for field, field_type, mj_type in ( + (mjm.opt.integrator, types.IntegratorType, mujoco.mjtIntegrator), + (mjm.opt.cone, types.ConeType, mujoco.mjtCone), + (mjm.opt.solver, types.SolverType, mujoco.mjtSolver), ): if field not in set(field_type): - raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.") + raise NotImplementedError(f"{mj_type(field).name} is unsupported.") # opt: check supported features in scalar flag types - for field, field_type in ( - (mjm.opt.disableflags, types.DisableBit), - (mjm.opt.enableflags, types.EnableBit), + for field, field_type, mj_type in ( + (mjm.opt.disableflags, types.DisableBit, mujoco.mjtDisableBit), + (mjm.opt.enableflags, types.EnableBit, mujoco.mjtEnableBit), ): - if field & ~np.bitwise_or.reduce(field_type): - raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.") + unsupported = field & ~np.bitwise_or.reduce(field_type) + if unsupported: + raise NotImplementedError(f"{mj_type(unsupported).name} is unsupported.") if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any(): raise NotImplementedError("Flex collisions are not implemented.") @@ -153,6 +176,21 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: if not_implemented(objtype, objid, types.GeomType.BOX) and not_implemented(reftype, refid, types.GeomType.BOX): raise NotImplementedError(f"Collision sensors with box-box collisions are not implemented.") + def _check_friction(name: str, id_: int, condim: int, friction, checks): + for min_condim, indices in checks: + if condim >= min_condim: + for idx in indices: + if friction[idx] < types.MJ_MINMU: + warnings.warn( + f"{name} {id_}: friction[{idx}] ({friction[idx]}) < MJ_MINMU ({types.MJ_MINMU}) with condim={condim} may cause NaN" + ) + + for geomid in range(mjm.ngeom): + _check_friction("geom", geomid, mjm.geom_condim[geomid], mjm.geom_friction[geomid], [(3, [0]), (4, [1]), (6, [2])]) + + for pairid in range(mjm.npair): + _check_friction("pair", pairid, mjm.pair_dim[pairid], mjm.pair_friction[pairid], [(3, [0]), (4, [1, 2]), (6, [3, 4])]) + # create opt opt_kwargs = {f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)} if hasattr(mjm.opt, "impratio"): @@ -856,16 +894,22 @@ def put_data( 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)) d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) - flexedge_J = np.zeros((mjm.nflexedge, mjm.nv)) - mujoco.mju_sparse2dense( - flexedge_J, mjd.flexedge_J.reshape(-1), mjm.flexedge_J_rownnz, mjm.flexedge_J_rowadr, mjm.flexedge_J_colind.reshape(-1) - ) - d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float) else: 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) - flexedge_J = mjd.flexedge_J.reshape((mjm.nflexedge, mjm.nv)) - d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float) + + flexedge_J = np.zeros((mjm.nflexedge, mjm.nv)) + if mjd.flexedge_J.size: + # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml + if BLEEDING_EDGE_MUJOCO: + mujoco.mju_sparse2dense( + flexedge_J, mjd.flexedge_J.reshape(-1), mjm.flexedge_J_rownnz, mjm.flexedge_J_rowadr, mjm.flexedge_J_colind.reshape(-1) + ) + else: + mujoco.mju_sparse2dense( + flexedge_J, mjd.flexedge_J.reshape(-1), mjd.flexedge_J_rownnz, mjd.flexedge_J_rowadr, mjd.flexedge_J_colind.reshape(-1) + ) + d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float) # TODO(taylorhowell): sparse actuator_moment actuator_moment = np.zeros((mjm.nu, mjm.nv)) @@ -980,16 +1024,31 @@ def get_data_into( result.cdof[:] = d.cdof.numpy()[world_id] result.cinert[:] = d.cinert.numpy()[world_id] result.flexvert_xpos[:] = d.flexvert_xpos.numpy()[world_id] - result.flexedge_J[:] = d.flexedge_J.numpy()[world_id] + flexedge_J = d.flexedge_J.numpy()[world_id] + if result.flexedge_J.size: + # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml + if BLEEDING_EDGE_MUJOCO: + mujoco.mju_dense2sparse( + result.flexedge_J.reshape(-1), + flexedge_J, + mjm.flexedge_J_rownnz, + mjm.flexedge_J_rowadr, + mjm.flexedge_J_colind.reshape(-1), + ) + else: + mujoco.mju_dense2sparse( + result.flexedge_J.reshape(-1), + flexedge_J, + result.flexedge_J_rownnz, + result.flexedge_J_rowadr, + result.flexedge_J_colind.reshape(-1), + ) result.flexedge_length[:] = d.flexedge_length.numpy()[world_id] result.flexedge_velocity[:] = d.flexedge_velocity.numpy()[world_id] result.actuator_length[:] = d.actuator_length.numpy()[world_id] + actuator_moment = d.actuator_moment.numpy()[world_id] mujoco.mju_dense2sparse( - result.actuator_moment, - d.actuator_moment.numpy()[world_id], - result.moment_rownnz, - result.moment_rowadr, - result.moment_colind, + result.actuator_moment, actuator_moment, result.moment_rownnz, result.moment_rowadr, result.moment_colind ) result.crb[:] = d.crb.numpy()[world_id] result.qLDiagInv[:] = d.qLDiagInv.numpy()[world_id] 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 3ee4ffb3..ed59f5ad 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -595,15 +595,15 @@ def _flex_elasticity( nedge = nvert * (nvert - 1) / 2 edges = wp.where( dim == 3, - wp.types.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int), - wp.types.matrix(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int), + wp.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int), + wp.matrix(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int), ) if timestep > 0.0 and not dsbl_damper: kD = flex_damping[f] / timestep else: kD = 0.0 - gradient = wp.types.matrix(0.0, shape=(6, 6)) + gradient = wp.matrix(0.0, shape=(6, 6)) for e in range(nedge): vert0 = flex_elem[(dim + 1) * elemid + edges[e, 0]] vert1 = flex_elem[(dim + 1) * elemid + edges[e, 1]] @@ -622,7 +622,7 @@ def _flex_elasticity( previous = deformed - vel * timestep elongation[e] = deformed * deformed - reference * reference + (deformed * deformed - previous * previous) * kD - metric = wp.types.matrix(0.0, shape=(6, 6)) + metric = wp.matrix(0.0, shape=(6, 6)) id = int(0) for ed1 in range(nedge): for ed2 in range(ed1, nedge): @@ -630,7 +630,7 @@ def _flex_elasticity( metric[ed2, ed1] = flex_stiffness[elemid, id] id += 1 - force = wp.types.matrix(0.0, shape=(6, 3)) + force = wp.matrix(0.0, shape=(6, 3)) for ed1 in range(nedge): for ed2 in range(nedge): for i in range(2): @@ -684,7 +684,7 @@ def _flex_bending( flex_vertadr[f] + flex_edgeflap[edgeid][1], ) - frc = wp.types.matrix(0.0, shape=(4, 3)) + frc = wp.matrix(0.0, shape=(4, 3)) if flex_bending[edgeid, 16]: v0 = flexvert_xpos_in[worldid, v[0]] v1 = flexvert_xpos_in[worldid, v[1]] @@ -695,7 +695,7 @@ def _flex_bending( frc[3] = wp.cross(v1 - v0, v2 - v0) frc[0] = -(frc[1] + frc[2] + frc[3]) - force = wp.types.matrix(0.0, shape=(nvert, 3)) + force = wp.matrix(0.0, shape=(nvert, 3)) for i in range(nvert): for x in range(3): for j in range(nvert): @@ -786,24 +786,24 @@ def passive(m: Model, d: Data): ], outputs=[d.qfrc_spring], ) - wp.launch( - _flex_bending, - dim=(d.nworld, m.nflexedge), - inputs=[ - m.nflex, - m.body_dofadr, - m.flex_dim, - m.flex_vertadr, - m.flex_edgeadr, - m.flex_edgenum, - m.flex_vertbodyid, - m.flex_edge, - m.flex_edgeflap, - m.flex_bending, - d.flexvert_xpos, - ], - outputs=[d.qfrc_spring], - ) + wp.launch( + _flex_bending, + dim=(d.nworld, m.nflexedge), + inputs=[ + m.nflex, + m.body_dofadr, + m.flex_dim, + m.flex_vertadr, + m.flex_edgeadr, + m.flex_edgenum, + m.flex_vertbodyid, + m.flex_edge, + m.flex_edgeflap, + m.flex_bending, + d.flexvert_xpos, + ], + outputs=[d.qfrc_spring], + ) gravcomp = m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY) 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 41ff3008..76b29385 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -572,7 +572,7 @@ def _sensor_pos( elif sensortype == SensorType.FRAMEZAXIS: axis = 2 vec3 = _frame_axis( - ximat_in, xmat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis + xmat_in, ximat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: 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 43fb8c37..429c7a73 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -247,32 +247,6 @@ def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): apply_ft(m, d, d.xfrc_applied, qfrc, True) -@wp.func -def all_same(v0: wp.vec3, v1: wp.vec3) -> wp.bool: - dx = abs(v0[0] - v1[0]) - dy = abs(v0[1] - v1[1]) - dz = abs(v0[2] - v1[2]) - - return ( - (dx <= 1.0e-9 or dx <= max(abs(v0[0]), abs(v1[0])) * 1.0e-9) - and (dy <= 1.0e-9 or dy <= max(abs(v0[1]), abs(v1[1])) * 1.0e-9) - and (dz <= 1.0e-9 or dz <= max(abs(v0[2]), abs(v1[2])) * 1.0e-9) - ) - - -@wp.func -def any_different(v0: wp.vec3, v1: wp.vec3) -> wp.bool: - dx = abs(v0[0] - v1[0]) - dy = abs(v0[1] - v1[1]) - dz = abs(v0[2] - v1[2]) - - return ( - (dx > 1.0e-9 and dx > max(abs(v0[0]), abs(v1[0])) * 1.0e-9) - or (dy > 1.0e-9 and dy > max(abs(v0[1]), abs(v1[1])) * 1.0e-9) - or (dz > 1.0e-9 and dz > max(abs(v0[2]), abs(v1[2])) * 1.0e-9) - ) - - @wp.func def _decode_pyramid( njmax_in: int, pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int 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 6ff16d70..99cdbb97 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -25,7 +25,7 @@ MJ_MAXIMP = mujoco.mjMAXIMP # maximum constraint impedance MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR MJ_MINMU = mujoco.mjMINMU # minimum friction # maximum size (by number of edges) of an horizon in EPA algorithm -MJ_MAX_EPAHORIZON = 12 +MJ_MAX_EPAHORIZON = 24 # maximum average number of trianglarfaces EPA can insert at each iteration MJ_MAX_EPAFACES = 5 @@ -173,11 +173,13 @@ class EnableBit(enum.IntFlag): Attributes: ENERGY: energy computation INVDISCRETE: discrete-time inverse dynamics + MULTICCD: multiple contacts with CCD """ ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE - # unsupported: OVERRIDE, FWDINV, ISLAND, MULTICCD + MULTICCD = mujoco.mjtEnableBit.mjENBL_MULTICCD + # unsupported: OVERRIDE, FWDINV, ISLAND class TrnType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py index 1544ec63..0301a386 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py @@ -18,8 +18,6 @@ import inspect from typing import Callable, Optional import warp as wp -from warp._src.context import Module -from warp._src.context import get_module _STACK = None @@ -127,7 +125,7 @@ def nested_kernel( f: Optional[Callable] = None, *, enable_backward: Optional[bool] = None, - module: Optional[Module] = None, + module: Optional[wp.Module] = None, ): """Decorator to register a Warp kernel from a Python function. @@ -166,7 +164,7 @@ def nested_kernel( Args: f: The function to be registered as a kernel. enable_backward: If False, the backward pass will not be generated. - module: The :class:`warp.context.Module` to which the kernel belongs. Alternatively, + module: The :class:`warp.Module` to which the kernel belongs. Alternatively, if a string `"unique"` is provided, the kernel is assigned to a new module named after the kernel name and hash. If None, the module is inferred from the function's module. @@ -182,7 +180,7 @@ def nested_kernel( qualname = func.__qualname__ parts = [part for part in qualname.split(".") if part != ""] outer_functions = parts[:-1] - module_name = get_module(".".join([func.__module__] + outer_functions)) + module_name = wp.get_module(".".join([func.__module__] + outer_functions)) else: module_name = module @@ -220,8 +218,7 @@ def cache_kernel(func): def check_toolkit_driver(): - if wp._src.context.runtime is None: - wp._src.context.init() + wp.init() if wp.get_device().is_cuda: - if wp._src.context.runtime.toolkit_version < (12, 4) or wp._src.context.runtime.driver_version < (12, 4): + if not wp.is_conditional_graph_supported(): RuntimeError("Minimum supported CUDA version: 12.4.") diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml index 7006d4c8..94d6f7ea 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name="mujoco-warp" -version = "0.0.1" +version = "0.0.2" # TODO(team): create a distribution list authors = [ {name = "Newton Developers", email = "mujoco@deepmind.com"}, @@ -27,9 +27,9 @@ requires-python = ">=3.10" dependencies = [ "absl-py", "etils[epath]", - "mujoco>=3.3.7", + "mujoco>=3.4.0", "numpy", - "warp-lang>=1.9.1", + "warp-lang>=1.11.0", ] [[tool.uv.index]] @@ -48,7 +48,6 @@ mujoco = {index = "mujoco"} [project.optional-dependencies] dev = [ - "asv", "pre-commit", "pytest", "pytest-xdist", @@ -56,7 +55,7 @@ dev = [ "pygls>=1.0.0,<2.0.0", "lsprotocol>=2023.0.1,<2024.0.0", "mujoco>=3.3.7.dev0", - "warp-lang>=1.9.1.dev0", + "warp-lang>=1.11.0.dev0", ] # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp cpu = [ @@ -113,6 +112,10 @@ convention = "google" docstring-code-format = true docstring-code-line-length = 100 +[tool.pytest.ini_options] +testpaths = ["mujoco_warp"] +norecursedirs = ["benchmarks", "contrib"] + [tool.setuptools] package-data = { "mujoco_warp" = ["test_data/**"] } diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py index 7d1b9afd..3cdb2c7f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py @@ -121,7 +121,6 @@ def _main(argv: Sequence[str]) -> None: mujoco.mj_resetDataKeyframe(mjm, mjd, keys[0]) elif mjm.nkey > 0 and _KEYFRAME.value > -1: mujoco.mj_resetDataKeyframe(mjm, mjd, _KEYFRAME.value) - mujoco.mj_forward(mjm, mjd) if _ENGINE.value == EngineOptions.C: override_model(mjm, _OVERRIDE.value) diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 446083ed..581888ec 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -42,6 +42,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _collision_shim( # Model @@ -108,6 +109,7 @@ def _collision_shim( opt__ccd_iterations: int, opt__ccd_tolerance: wp.array(dtype=float), opt__disableflags: int, + opt__enableflags: int, opt__sdf_initpoints: int, opt__sdf_iterations: int, # Data @@ -190,6 +192,7 @@ def _collision_shim( _m.opt.ccd_iterations = opt__ccd_iterations _m.opt.ccd_tolerance = opt__ccd_tolerance _m.opt.disableflags = opt__disableflags + _m.opt.enableflags = opt__enableflags _m.opt.sdf_initpoints = opt__sdf_initpoints _m.opt.sdf_iterations = opt__sdf_iterations _m.pair_dim = pair_dim @@ -341,6 +344,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data): m.opt._impl.ccd_iterations, m.opt._impl.ccd_tolerance, m.opt.disableflags, + m.opt.enableflags, m.opt._impl.sdf_initpoints, m.opt._impl.sdf_iterations, d._impl.naconmax, diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index b2734529..0e92a3c8 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/mjx/mujoco/mjx/warp/ffi.py @@ -150,7 +150,7 @@ def _format_arg(arg: Any, name: str, annotation: Any, verbose: bool): for i in range(len(arg)) ) - if not isinstance(annotation, wp.types.array): + if not isinstance(annotation, wp.array): if verbose: print(f'Skipping {name}: {arg}') return arg