diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt index a540879c..b2e81f22 100644 --- a/mjx/cuda_requirements.txt +++ b/mjx/cuda_requirements.txt @@ -16,8 +16,12 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \ jax-cuda12-pjrt==0.4.30; python_version == '3.9' \ --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \ --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb -warp-lang==1.10.0 \ - --hash=sha256:428a6388ba8c9b3ded973226ddb5be59b16e3b3f28ac939a5036ad2d7cdc79ed \ - --hash=sha256:1bdff31e170b00c89fb9d8b647e906fefdcdcf8741925126d1c0be42783174fa \ - --hash=sha256:4aa8eb63cae5ee0d6dbdbdfc305d124140ec975475c97d4458f412396fb39eab \ - --hash=sha256:81f73055e76a6a3f2284cf2b5fe542a341156861101df9f291bd5b59925ff6e5 +warp-lang==1.10.1 \ + --hash=sha256:0c6f44d4136cfc86316f5d35883a863b4dca2bc318331597e755ea872db0cd97 \ + --hash=sha256:2884b642f16b07b930b3605193c1b97d183eef80c5b0083d3a17473aee92138e \ + --hash=sha256:7c8f839a01042677d1f417d31abe9304f9319fca7041c2ed82937ed511eb2362 \ + --hash=sha256:7addd14a913b50c406ba3b1abebf3a208d01591a1c382da799cff2806d234c5a \ + --hash=sha256:5f5ce0147d48f86ccab66b834aa54c8adedbbb2ba86e45b2b904b182fa963509 \ + --hash=sha256:2068e877dd2972d841c1eb63f98eadce326e77ba58f09e8474d0847276f47633 \ + --hash=sha256:96d24694333e30eac888b2c2962bc5ca4b2369f56707e4c54a465fa37b7ae007 \ + --hash=sha256:b77ffe935e06cb57e76f6d4dd8bae7d59294ce63977ea8a0a4fd8e30d77c0a7f diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py index 4b9c45da..744894df 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py @@ -56,6 +56,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.smooth import com_pos as com_pos from mujoco.mjx.third_party.mujoco_warp._src.smooth import com_vel as com_vel from mujoco.mjx.third_party.mujoco_warp._src.smooth import crb as crb from mujoco.mjx.third_party.mujoco_warp._src.smooth import factor_m as factor_m +from mujoco.mjx.third_party.mujoco_warp._src.smooth import flex as flex from mujoco.mjx.third_party.mujoco_warp._src.smooth import kinematics as kinematics from mujoco.mjx.third_party.mujoco_warp._src.smooth import rne as rne from mujoco.mjx.third_party.mujoco_warp._src.smooth import rne_postconstraint as rne_postconstraint 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 32ea187c..f72ceea4 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py @@ -174,6 +174,7 @@ class BenchmarkSuite: batch_size = -1 nconmax = -1 njmax = -1 + nstep = 1000 param_names = ("function",) params = ( "jit_duration", @@ -239,7 +240,7 @@ class BenchmarkSuite: d = io.put_data(mjm, mjd, self.batch_size, self.nconmax, self.njmax) free_after = wp.get_device().free_memory - jit_duration, _, trace, _, _, solver_niter, _ = benchmark(forward.step, m, d, 1000, ctrls, True, False, True) + 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), 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 4a25dd21..16d67051 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 @@ -23,9 +23,8 @@ def create_blocked_cholesky_func(block_size: int): @wp.func def blocked_cholesky_func( # In: - tid_block: int, A: wp.array(dtype=float, ndim=2), - active_matrix_size: int, + matrix_size: int, # Out: L: wp.array(dtype=float, ndim=2), ): @@ -33,89 +32,57 @@ def create_blocked_cholesky_func(block_size: int): It returns a lower-triangular matrix L such that A = L L^T. """ - num_threads_per_block = wp.block_dim() - - # Round up active_matrix_size to next multiple of block_size - n = ((active_matrix_size + block_size - 1) // block_size) * block_size + # TODO(team): remove conditional after mjwarp relies on >= 1.11 + bleeding_edge_warp = wp.static(wp.__version__ >= "1.11") # Process the matrix in blocks along its leading dimension. - for k in range(0, n, block_size): + for k in range(0, matrix_size, block_size): end = k + block_size # Load current diagonal block A[k:end, k:end] # and update with contributions from previously computed blocks. A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared") - # The following if pads the matrix if it is not divisible by block_size - if k + block_size > active_matrix_size or k + block_size > active_matrix_size: - num_tile_elements = block_size * block_size - num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block - for i in range(num_iterations): - linear_index = tid_block + i * num_threads_per_block - linear_index = linear_index % num_tile_elements - row = linear_index // block_size - col = linear_index % block_size - value = A_kk_tile[row, col] - if k + row >= active_matrix_size or k + col >= active_matrix_size: - value = wp.where(row == col, float(1), float(0)) - A_kk_tile[row, col] = value - - if k > 0: - for j in range(0, k, block_size): - L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j)) - L_block_T = wp.tile_transpose(L_block) - L_L_T_block = wp.tile_matmul(L_block, L_block_T) - A_kk_tile -= L_L_T_block + 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)) # Compute the Cholesky factorization for the block L_kk_tile = wp.tile_cholesky(A_kk_tile) wp.tile_store(L, L_kk_tile, offset=(k, k)) # Process the blocks below the current block - for i in range(end, n, block_size): + for i in range(end, matrix_size, block_size): A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared") - # The following if pads the matrix if it is not divisible by block_size - if i + block_size > active_matrix_size or k + block_size > active_matrix_size: - num_tile_elements = block_size * block_size - num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block - for ii in range(num_iterations): - linear_index = tid_block + ii * num_threads_per_block - linear_index = linear_index % num_tile_elements - row = linear_index // block_size - col = linear_index % block_size - value = A_ik_tile[row, col] - if i + row >= active_matrix_size or k + col >= active_matrix_size: - value = wp.where(i + row == k + col, float(1), float(0)) - A_ik_tile[row, col] = value + 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)) - if k > 0: - for j in range(0, k, block_size): - L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j)) - L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j)) - L_T_tile = wp.tile_transpose(L_2_tile) - L_L_T_tile = wp.tile_matmul(L_tile, L_T_tile) - A_ik_tile -= L_L_T_tile - - t = wp.tile_transpose(A_ik_tile) - tmp = wp.tile_lower_solve(L_kk_tile, t) - sol_tile = wp.tile_transpose(tmp) - - wp.tile_store(L, sol_tile, offset=(i, k)) + 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_store(L, A_ik_tile, offset=(i, k)) return blocked_cholesky_func @lru_cache(maxsize=None) -def create_blocked_cholesky_solve_func(block_size: int): +def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int): @wp.func def blocked_cholesky_solve_func( # In: - tid_block: int, L: wp.array(dtype=float, ndim=2), b: wp.array(dtype=float, ndim=2), - tmp: wp.array(dtype=float, ndim=2), - active_matrix_size: int, + matrix_size: int, # Out: x: wp.array(dtype=float, ndim=2), ): @@ -124,94 +91,45 @@ def create_blocked_cholesky_solve_func(block_size: int): Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward substitution. """ - num_threads_per_block = wp.block_dim() - - # Round up active_matrix_size to next multiple of block_size - n = ((active_matrix_size + block_size - 1) // block_size) * block_size + # TODO(team): remove conditional after mjwarp relies on >= 1.11 + bleeding_edge_warp = wp.static(wp.__version__ >= "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 - for i in range(0, n, block_size): - i_end = i + block_size - rhs_tile = wp.tile_load(b, shape=(block_size, 1), offset=(i, 0)) - if i > 0: - for j in range(0, i, block_size): - L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j)) - y_block = wp.tile_load(tmp, shape=(block_size, 1), offset=(j, 0)) - Ly_block = wp.tile_matmul(L_block, y_block) - rhs_tile -= Ly_block - L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i)) + for i in range(0, matrix_size, block_size): + rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0)) + 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) - # The following if pads the matrix if it is not divisible by block_size - if i + block_size > active_matrix_size: - num_tile_elements = block_size * block_size - num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block - - for ii in range(num_iterations): - linear_index = tid_block + ii * num_threads_per_block - linear_index = linear_index % num_tile_elements - row = linear_index // block_size - col = linear_index % block_size - value = L_tile[row, col] - if i + row >= active_matrix_size or i + col >= active_matrix_size: - value = wp.where(row == col, float(1), float(0)) - L_tile[row, col] = value - - # Handle rhs - num_tile_elements = block_size - num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block - - for ii in range(num_iterations): - linear_index = tid_block + ii * num_threads_per_block - linear_index = linear_index % num_tile_elements - value = rhs_tile[linear_index, 0] - if i + linear_index >= active_matrix_size: - value = float(0) - rhs_tile[linear_index, 0] = value - - y_tile = wp.tile_lower_solve(L_tile, rhs_tile) - wp.tile_store(tmp, y_tile, offset=(i, 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)) # Backward substitution: solve L^T x = y - for i in range(n - block_size, -1, -block_size): + for i in range(matrix_size - block_size, -1, -block_size): i_end = i + block_size - rhs_tile = wp.tile_load(tmp, shape=(block_size, 1), offset=(i, 0)) - if i_end < n: - for j in range(i_end, n, block_size): - L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i)) - L_T_tile = wp.tile_transpose(L_tile) - x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, 0)) - L_T_x_tile = wp.tile_matmul(L_T_tile, x_tile) - rhs_tile -= L_T_x_tile - L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i)) + tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0)) + 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) + L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared") - # The following if pads the matrix if it is not divisible by block_size - if i + block_size > active_matrix_size: - num_tile_elements = block_size * block_size - num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block - - for ii in range(num_iterations): - linear_index = tid_block + ii * num_threads_per_block - linear_index = linear_index % num_tile_elements - row = linear_index // block_size - col = linear_index % block_size - value = L_tile[row, col] - if i + row >= active_matrix_size or i + col >= active_matrix_size: - value = wp.where(row == col, float(1), float(0)) - L_tile[row, col] = value - - # Handle rhs - num_tile_elements = block_size - num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block - - for ii in range(num_iterations): - linear_index = tid_block + ii * num_threads_per_block - linear_index = linear_index % num_tile_elements - value = rhs_tile[linear_index, 0] - if i + linear_index >= active_matrix_size: - value = float(0) - rhs_tile[linear_index, 0] = value - - x_tile = wp.tile_upper_solve(wp.tile_transpose(L_tile), rhs_tile) - wp.tile_store(x, x_tile, offset=(i, 0)) + 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_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 d3772f89..de3f485b 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 @@ -195,8 +195,8 @@ def _hfield_filter( def ccd_kernel_builder( geomtype1: int, geomtype2: int, - ccd_iterations: int, - is_hfield: bool, + gjk_iterations: int, + epa_iterations: int, use_multiccd: bool, ): @wp.func @@ -274,7 +274,8 @@ def ccd_kernel_builder( dist, ncontact, w1, w2, idx = ccd( opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], cutoff, - ccd_iterations, + gjk_iterations, + epa_iterations, geom1, geom2, geomtype1, @@ -484,7 +485,7 @@ def ccd_kernel_builder( worldid = collision_worldid_in[tid] # height field filter - if wp.static(is_hfield): + if wp.static(geomtype1 == GeomType.HFIELD.value): 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 ) @@ -538,7 +539,7 @@ def ccd_kernel_builder( ) # see MuJoCo mjc_ConvexHField - if wp.static(is_hfield): + if wp.static(geomtype1 == GeomType.HFIELD.value): geom1_dataid = geom_dataid[g1] # height field subgrid @@ -645,7 +646,8 @@ def ccd_kernel_builder( dist, ncontact, w1, w2, idx = ccd( opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], 0.0, - ccd_iterations, + gjk_iterations, + epa_iterations, geom1, geom2, geomtype1, @@ -967,34 +969,41 @@ def convex_narrowphase(m: Model, d: Data): kernel for each type of convex collision pair present in the model, avoiding unnecessary computations for non-existent pair types. """ - if not any(m.geom_pair_type_count[upper_trid_index(len(GeomType), g[0].value, g[1].value)] for g in _CONVEX_COLLISION_PAIRS): + + def _pair_count(p1: int, p2: int) -> int: + return m.geom_pair_type_count[upper_trid_index(len(GeomType), p1, p2)] + + # no convex collisions, early return + if not any(_pair_count(g[0].value, g[1].value) for g in _CONVEX_COLLISION_PAIRS): return + epa_iterations = m.opt.ccd_iterations + # set to true to enable multiccd use_multiccd = False 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 + m.opt.ccd_iterations), dtype=wp.vec3) + 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 + m.opt.ccd_iterations), dtype=wp.vec3) + epa_vert1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3) # epa_vert2: vertices in EPA polytope in geom 2 space - epa_vert2 = wp.empty(shape=(d.naconmax, 5 + m.opt.ccd_iterations), dtype=wp.vec3) + epa_vert2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3) # epa_vert_index1: vertex indices in EPA polytope for geom 1 - epa_vert_index1 = wp.empty(shape=(d.naconmax, 5 + m.opt.ccd_iterations), dtype=int) + epa_vert_index1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int) # epa_vert_index2: vertex indices in EPA polytope for geom 2 (naconmax, 5 + CCDiter) - epa_vert_index2 = wp.empty(shape=(d.naconmax, 5 + m.opt.ccd_iterations), dtype=int) + 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 * m.opt.ccd_iterations), dtype=wp.vec3i) + epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3i) # epa_pr: projection of origin on polytope faces - epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * m.opt.ccd_iterations), dtype=wp.vec3) + 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 * m.opt.ccd_iterations), dtype=float) + 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 * m.opt.ccd_iterations), dtype=int) + 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 * m.opt.ccd_iterations), dtype=int) + 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) # multiccd_polygon: clipped contact surface @@ -1023,9 +1032,9 @@ def convex_narrowphase(m: Model, d: Data): for geom_pair in _CONVEX_COLLISION_PAIRS: g1 = geom_pair[0].value g2 = geom_pair[1].value - if m.geom_pair_type_count[upper_trid_index(len(GeomType), g1, g2)]: + if _pair_count(g1, g2): wp.launch( - ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, g1 == GeomType.HFIELD, use_multiccd), + ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd), dim=d.naconmax, inputs=[ m.opt.ccd_tolerance, 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 8d59525a..bc2e2513 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 @@ -18,7 +18,6 @@ from typing import Tuple 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 @@ -28,7 +27,8 @@ wp.set_module_options({"enable_backward": False}) FLOAT_MIN = -1e30 FLOAT_MAX = 1e30 -MJ_MINVAL2 = MJ_MINVAL * MJ_MINVAL +MINVAL = 1e-15 +MIN_DIST = 1e-10 # TODO(kbayes): write out formulas to derive these constants FACE_TOL = 0.99999872 @@ -122,7 +122,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: res = wp.vec3(0.0, 0.0, 0.0) # set result in XY plane: support on circle d = wp.sqrt(local_dir[0] * local_dir[0] + local_dir[1] * local_dir[1]) - if d > MJ_MINVAL: + if d > MINVAL: scl = geom.size[0] / d res[0] = local_dir[0] * scl res[1] = local_dir[1] * scl @@ -173,6 +173,8 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: sp.point = geom.rot @ sp.point + geom.pos 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) for i in range(6): vert = geom.hfprism[i] dist = wp.dot(vert, local_dir) @@ -241,7 +243,7 @@ def _linear_combine(n: int, coefs: wp.vec4, mat: mat43) -> wp.vec3: @wp.func def _almost_equal(v1: wp.vec3, v2: wp.vec3) -> bool: - return wp.abs(v1[0] - v2[0]) < MJ_MINVAL and wp.abs(v1[1] - v2[1]) < MJ_MINVAL and wp.abs(v1[2] - v2[2]) < MJ_MINVAL + return wp.abs(v1[0] - v2[0]) < MINVAL and wp.abs(v1[1] - v2[1]) < MINVAL and wp.abs(v1[2] - v2[2]) < MINVAL @wp.func @@ -291,7 +293,7 @@ def _project_origin_plane(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> Tuple[wp.vec nn = wp.dot(n, n) if nn == 0: return z, 1 - if nv != 0 and nn > MJ_MINVAL: + if nv != 0 and nn > MINVAL: v = (nv / nn) * n return v, 0 @@ -301,7 +303,7 @@ def _project_origin_plane(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> Tuple[wp.vec nn = wp.dot(n, n) if nn == 0: return z, 1 - if nv != 0 and nn > MJ_MINVAL: + if nv != 0 and nn > MINVAL: v = (nv / nn) * n return v, 0 @@ -579,13 +581,14 @@ def gjk( simplex_index1 = wp.vec4i() simplex_index2 = wp.vec4i() n = int(0) + cnt = int(1) coordinates = wp.vec4() # barycentric coordinates epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance) # set initial guess x_k = x1_0 - x2_0 - for k in range(gjk_iterations): + for _ in range(gjk_iterations): xnorm = wp.dot(x_k, x_k) # TODO(kbayes): determine new constant here if xnorm < 1e-12: @@ -663,6 +666,11 @@ def gjk( if n == 4: break + cnt += 1 + + if cnt == gjk_iterations: + wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations) + result = GJKResult() # compute the approximate witness points @@ -751,7 +759,7 @@ def _tri_point_intersect(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, p: wp.vec3) -> b pr[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3 pr[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3 pr[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3 - return wp.norm_l2(pr - p) < MJ_MINVAL + return wp.norm_l2(pr - p) < MINVAL @wp.func @@ -867,36 +875,66 @@ def _delete_face(pt: Polytope, face_id: int) -> int: @wp.func -def _epa_witness(pt: Polytope, face_idx: int) -> Tuple[wp.vec3, wp.vec3]: +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] + # compute affine coordinates for witness points on plane defined by face - v1 = pt.vert[pt.face[face_idx][0]] - v2 = pt.vert[pt.face[face_idx][1]] - v3 = pt.vert[pt.face[face_idx][2]] + v1 = pt.vert[face[0]] + v2 = pt.vert[face[1]] + v3 = pt.vert[face[2]] coordinates = _tri_affine_coord(v1, v2, v3, pt.face_pr[face_idx]) l1 = coordinates[0] l2 = coordinates[1] l3 = coordinates[2] - # face on geom 1 - v1 = pt.vert1[pt.face[face_idx][0]] - v2 = pt.vert1[pt.face[face_idx][1]] - v3 = pt.vert1[pt.face[face_idx][2]] - x1 = wp.vec3() - x1[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3 - x1[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3 - x1[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3 - # face on geom 2 - v1 = pt.vert2[pt.face[face_idx][0]] - v2 = pt.vert2[pt.face[face_idx][1]] - v3 = pt.vert2[pt.face[face_idx][2]] + v1 = pt.vert2[face[0]] + v2 = pt.vert2[face[1]] + v3 = pt.vert2[face[2]] x2 = wp.vec3() x2[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3 x2[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3 x2[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3 - return x1, x2 + # correct witness points for hfield geoms + i1 = pt.vert_index1[face[0]] + i2 = pt.vert_index1[face[1]] + 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] + a = geom1.hfprism[3] + b = geom1.hfprism[4] + c = geom1.hfprism[5] + x2 = wp.normalize(x2) + + # TODO(kbayes): Support cases where geom2 is larger than the height field + sp = _support(geom2, geomtype2, x2) + x2 = sp.point + + coordinates2 = _tri_affine_coord(a, b, c, x2) + if coordinates2[0] > 0 and coordinates2[1] > 0 and coordinates2[2] > 0: + x1 = coordinates[0] * a + coordinates[1] * b + coordinates[2] * c + else: + p = c + p = wp.where(coordinates[1] > 0, b, p) + p = wp.where(coordinates[0] > 0, a, p) + x1 = x2 - wp.dot(x2 - p, n) * n + return x1, x2, -wp.norm_l2(x1 - x2) + + # face on geom 1 + v1 = pt.vert1[face[0]] + v2 = pt.vert1[face[1]] + v3 = pt.vert1[face[2]] + x1 = wp.vec3() + x1[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3 + x1[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3 + x1[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3 + + return x1, x2, -wp.sqrt(pt.face_norm2[face_idx]) @wp.func @@ -926,7 +964,7 @@ def _polytope2( index = i # cross product with best coordinate axis - e = wp.vec(0.0, 0.0, 0.0) + e = wp.vec3(0.0, 0.0, 0.0) e[index] = 1.0 d1 = wp.cross(e, diff) @@ -956,27 +994,27 @@ def _polytope2( _epa_support(pt, 4, geom1, geom2, geomtype1, geomtype2, d3 / wp.norm_l2(d3)) # build hexahedron - if _attach_face(pt, 0, 0, 2, 3) < MJ_MINVAL: + if _attach_face(pt, 0, 0, 2, 3) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 0, 2, 3) - if _attach_face(pt, 1, 0, 4, 2) < MJ_MINVAL2: + if _attach_face(pt, 1, 0, 4, 2) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 0, 4, 2) - if _attach_face(pt, 2, 0, 3, 4) < MJ_MINVAL2: + if _attach_face(pt, 2, 0, 3, 4) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 0, 3, 4) - if _attach_face(pt, 3, 1, 3, 2) < MJ_MINVAL2: + if _attach_face(pt, 3, 1, 3, 2) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 1, 3, 2) - if _attach_face(pt, 4, 1, 2, 4) < MJ_MINVAL2: + if _attach_face(pt, 4, 1, 2, 4) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 1, 2, 4) - if _attach_face(pt, 5, 1, 4, 3) < MJ_MINVAL2: + if _attach_face(pt, 5, 1, 4, 3) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 1, 4, 3) @@ -1016,7 +1054,7 @@ def _polytope3( """Create polytope for EPA given a 2-simplex from GJK.""" # get normals in both directions n = wp.cross(simplex[1] - simplex[0], simplex[2] - simplex[0]) - if wp.norm_l2(n) < MJ_MINVAL: + if wp.norm_l2(n) < MINVAL: pt.status = 2 return pt @@ -1066,22 +1104,22 @@ def _polytope3( return pt # create hexahedron for EPA - if _attach_face(pt, 0, 4, 0, 1) < MJ_MINVAL2: + if _attach_face(pt, 0, 4, 0, 1) < MIN_DIST: pt.status = 6 return pt - if _attach_face(pt, 1, 4, 2, 0) < MJ_MINVAL2: + if _attach_face(pt, 1, 4, 2, 0) < MIN_DIST: pt.status = 7 return pt - if _attach_face(pt, 2, 4, 1, 2) < MJ_MINVAL2: + if _attach_face(pt, 2, 4, 1, 2) < MIN_DIST: pt.status = 8 return pt - if _attach_face(pt, 3, 3, 1, 0) < MJ_MINVAL2: + if _attach_face(pt, 3, 3, 1, 0) < MIN_DIST: pt.status = 9 return pt - if _attach_face(pt, 4, 3, 0, 2) < MJ_MINVAL2: + if _attach_face(pt, 4, 3, 0, 2) < MIN_DIST: pt.status = 10 return pt - if _attach_face(pt, 5, 3, 2, 1) < MJ_MINVAL2: + if _attach_face(pt, 5, 3, 2, 1) < MIN_DIST: pt.status = 11 return pt @@ -1140,19 +1178,19 @@ def _polytope4( pt.vert_index2[3] = simplex_index2[3] # if the origin is on a face, replace the 3-simplex with a 2-simplex - if _attach_face(pt, 0, 0, 1, 2) < MJ_MINVAL2: + if _attach_face(pt, 0, 0, 1, 2) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 0, 1, 2) - if _attach_face(pt, 1, 0, 3, 1) < MJ_MINVAL2: + if _attach_face(pt, 1, 0, 3, 1) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 0, 3, 1) - if _attach_face(pt, 2, 0, 2, 3) < MJ_MINVAL2: + if _attach_face(pt, 2, 0, 2, 3) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 0, 2, 3) - if _attach_face(pt, 3, 3, 2, 1) < MJ_MINVAL2: + if _attach_face(pt, 3, 3, 2, 1) < MIN_DIST: pt.status = -1 return pt, _replace_simplex3(pt, 3, 2, 1) @@ -1177,6 +1215,7 @@ def _polytope4( def _epa( # In: tolerance: float, + gjk_iterations: int, epa_iterations: int, pt: Polytope, geom1: Geom, @@ -1191,9 +1230,10 @@ def _epa( idx = int(-1) pidx = int(-1) epsilon = wp.where(is_discrete, 1e-15, tolerance) + cnt = int(1) - for k in range(epa_iterations): - pidx = int(idx) + for _ in range(epa_iterations): + pidx = idx idx = int(-1) # find the face closest to the origin (lower bound for penetration depth) @@ -1283,11 +1323,15 @@ def _epa( # clear horizon pt.nhorizon = 0 + cnt += 1 + + if cnt == epa_iterations: + wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations) # return from valid face if idx > -1: - x1, x2 = _epa_witness(pt, idx) - return -wp.sqrt(pt.face_norm2[idx]), x1, x2, idx + x1, x2, dist = _epa_witness(pt, geom1, geom2, geomtype1, geomtype2, idx) + return dist, x1, x2, idx return 0.0, wp.vec3(), wp.vec3(), -1 @@ -1696,40 +1740,40 @@ def _box_edge_normals( def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int: # compute global coordinates of the box face and face normal if idx == 0: # right - face_out[0] = mat @ wp.vec(size[0], size[1], size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos - face_out[3] = mat @ wp.vec(size[0], -size[1], size[2]) + pos + face_out[0] = mat @ wp.vec3(size[0], size[1], size[2]) + pos + face_out[1] = mat @ wp.vec3(size[0], size[1], -size[2]) + pos + face_out[2] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos + face_out[3] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos return 4 if idx == 1: # left - face_out[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos - face_out[1] = mat @ wp.vec(-size[0], size[1], size[2]) + pos - face_out[2] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + face_out[0] = mat @ wp.vec3(-size[0], size[1], -size[2]) + pos + face_out[1] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos + face_out[2] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos + face_out[3] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos return 4 if idx == 2: # top - face_out[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], size[1], size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], size[1], size[2]) + pos + face_out[0] = mat @ wp.vec3(-size[0], size[1], -size[2]) + pos + face_out[1] = mat @ wp.vec3(size[0], size[1], -size[2]) + pos + face_out[2] = mat @ wp.vec3(size[0], size[1], size[2]) + pos + face_out[3] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos return 4 if idx == 3: # bottom - face_out[0] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], -size[1], size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + face_out[0] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos + face_out[1] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos + face_out[2] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos + face_out[3] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos return 4 if idx == 4: # front - face_out[0] = mat @ wp.vec(-size[0], size[1], size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], size[1], size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], -size[1], size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos + face_out[0] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos + face_out[1] = mat @ wp.vec3(size[0], size[1], size[2]) + pos + face_out[2] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos + face_out[3] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos return 4 if idx == 5: # back - face_out[0] = mat @ wp.vec(size[0], size[1], -size[2]) + pos - face_out[1] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos - face_out[2] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos - face_out[3] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos + face_out[0] = mat @ wp.vec3(size[0], size[1], -size[2]) + pos + face_out[1] = mat @ wp.vec3(-size[0], size[1], -size[2]) + pos + face_out[2] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos + face_out[3] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos return 4 return 0 @@ -1791,7 +1835,6 @@ def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> Tuple[fl @wp.func def _polygon_clip( # In: - prune: bool, plane_normal: wp.array(dtype=wp.vec3), plane_dist: wp.array(dtype=float), face1: wp.array(dtype=wp.vec3), @@ -1871,7 +1914,7 @@ def _polygon_clip( if npolygon < 1: return 0, witness1, witness2 - if prune and npolygon > 4: + if npolygon > 4: quad = _polygon_quad(polygon_out, npolygon) for i in range(4): witness2[i] = polygon_out[quad[i]] @@ -2102,27 +2145,20 @@ def multicontact( face2, ) - # TODO(kbayes): this approximates the contact direction, by scaling the face normal by the - # single contact direction's magnitude. This is effective, but polygonClip should compute - # this for each contact point. - approx_dir = wp.vec3() - # face1 is an edge; clip face1 against face2 if is_edge_contact_geom1: approx_dir = wp.norm_l2(dir) * n2[j] - return _polygon_clip(False, plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped) + return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped) # face2 is an edge; clip face2 against face1 if is_edge_contact_geom2: approx_dir = -wp.norm_l2(dir) * n1[j] - return _polygon_clip(False, plane_normal, plane_dist, face1, nface1, face2, nface2, n1[j], approx_dir, polygon, clipped) + return _polygon_clip(plane_normal, plane_dist, face1, nface1, face2, nface2, n1[j], approx_dir, polygon, clipped) # face-face collision approx_dir = wp.norm_l2(dir) * n2[j] - # don't prune box-box collisions (expect up to 8 contacts) - prune = not (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX) - return _polygon_clip(prune, plane_normal, plane_dist, face1, nface1, face2, nface2, n1[i], approx_dir, polygon, clipped) + return _polygon_clip(plane_normal, plane_dist, face1, nface1, face2, nface2, n1[i], approx_dir, polygon, clipped) @wp.func @@ -2142,7 +2178,8 @@ def ccd( # In: tolerance: float, cutoff: float, - ccd_iterations: int, + gjk_iterations: int, + epa_iterations: int, geom1: Geom, geom2: Geom, geomtype1: int, @@ -2176,7 +2213,8 @@ def ccd( geom1.margin = 0.0 geom1.size = wp.vec3(0.0, geom1.size[1], geom1.size[2]) - if geomtype2 == GeomType.SPHERE or geomtype2 == GeomType.CAPSULE: + # TODO(kbayes): support gjk margin trick with height fields + if geomtype1 != GeomType.HFIELD and (geomtype2 == GeomType.SPHERE or geomtype2 == GeomType.CAPSULE): size2 = geom2.size[0] full_margin2 = size2 + 0.5 * geom2.margin geom2.margin = 0.0 @@ -2185,7 +2223,7 @@ def ccd( # special handling for sphere and capsule (shrink to point and line respectively) if size1 + size2 > 0.0: cutoff += full_margin1 + full_margin2 - result = gjk(tolerance, ccd_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete) + result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete) # shallow penetration, inflate contact if result.dist > tolerance: @@ -2201,7 +2239,7 @@ def ccd( geom2.size = wp.vec3(size2, geom2.size[1], geom2.size[2]) cutoff -= full_margin1 + full_margin2 - result = gjk(tolerance, ccd_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete) + result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete) # no penetration depth to recover if result.dist > tolerance or result.dim < 2: @@ -2287,7 +2325,7 @@ def ccd( if pt.status: return result.dist, 1, result.x1, result.x2, -1 - dist, x1, x2, idx = _epa(tolerance, ccd_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete) + dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete) if idx == -1: return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1 return dist, 1, x1, x2, idx 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 4412c4cc..287becad 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 @@ -1035,7 +1035,7 @@ def plane_box_wrapper( dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size) frame = make_frame(normal) - for i in range(4): + for i in range(8): write_contact( naconmax_in, i, @@ -1562,16 +1562,9 @@ assert _check_primitive_collisions(), "_PRIMITIVE_COLLISIONS is in invalid order @cache_kernel -def _create_narrowphase_kernel(primitive_collisions_types, primitive_collisions_func): - # AD: no unique here: - # * we expect this generator to be called only once per model, so no repeated compilation - # * module="unique" is generating problems because it uses the function name as the key - # that in turn will cause multiple kernels to be generated with the same name - # this is mostly problematic in cases like the UTs where we don't clear the kernel cache - # between different tests. - - @nested_kernel(enable_backward=False) - def _primitive_narrowphase( +def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func): + @nested_kernel(module="unique", enable_backward=False) + def primitive_narrowphase( # Model: geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), @@ -1719,20 +1712,11 @@ def _create_narrowphase_kernel(primitive_collisions_types, primitive_collisions_ nacon_out, ) - return _primitive_narrowphase + return primitive_narrowphase -def _primitive_narrowphase_builder(m: Model): - _primitive_collisions_types = [] - _primitive_collisions_func = [] - - for types, func in _PRIMITIVE_COLLISIONS.items(): - idx = upper_trid_index(len(GeomType), types[0].value, types[1].value) - if m.geom_pair_type_count[idx] and types not in _primitive_collisions_types: - _primitive_collisions_types.append(types) - _primitive_collisions_func.append(func) - - return _create_narrowphase_kernel(_primitive_collisions_types, _primitive_collisions_func) +_PRIMITIVE_COLLISION_TYPES = [] +_PRIMITIVE_COLLISION_FUNC = [] @event_scope @@ -1751,10 +1735,17 @@ def primitive_narrowphase(m: Model, d: Data): the specific primitive collision types present in the model, avoiding unnecessary checks for non-existent collision pairs. """ - # we need to figure out how to keep the overhead of this small - not launching anything + # TODO(team): keep the overhead of this small - not launching anything # for pair types without collisions, as well as updating the launch dimensions. + + for types, func in _PRIMITIVE_COLLISIONS.items(): + idx = upper_trid_index(len(GeomType), types[0].value, types[1].value) + if m.geom_pair_type_count[idx] and types not in _PRIMITIVE_COLLISION_TYPES: + _PRIMITIVE_COLLISION_TYPES.append(types) + _PRIMITIVE_COLLISION_FUNC.append(func) + wp.launch( - _primitive_narrowphase_builder(m), + _primitive_narrowphase(_PRIMITIVE_COLLISION_TYPES, _PRIMITIVE_COLLISION_FUNC), dim=d.naconmax, inputs=[ m.geom_type, 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 6d93660f..04d2ea3c 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 @@ -316,7 +316,7 @@ def plane_box( box_pos: wp.vec3, box_rot: wp.mat33, box_size: wp.vec3, -) -> Tuple[wp.vec4, mat43f, wp.vec3]: +) -> Tuple[vec8f, mat83f, wp.vec3]: """Core contact geometry calculation for plane-box collision. Args: @@ -325,42 +325,36 @@ def plane_box( box_pos: Center position of the box. box_rot: Rotation matrix of the box. box_size: Half-extents of the box along each axis. + margin: Collision tolerance. Returns: - Vector of contact distances (wp.inf for unpopulated contacts). - Matrix of contact positions (one per row). - Contact normal vector. """ - corner = wp.vec3() center_dist = wp.dot(box_pos - plane_pos, plane_normal) - dist = wp.vec4(wp.inf) - pos = mat43f() + dist = vec8f(wp.inf) + pos = mat83f() # test all corners, pick bottom 4 - ncontact = int(0) for i in range(8): # get corner in local coordinates - corner.x = wp.where(i & 1, box_size.x, -box_size.x) - corner.y = wp.where(i & 2, box_size.y, -box_size.y) - corner.z = wp.where(i & 4, box_size.z, -box_size.z) + corner = wp.vec3( + wp.where(i & 1, box_size[0], -box_size[0]), + wp.where(i & 2, box_size[1], -box_size[1]), + wp.where(i & 4, box_size[2], -box_size[2]), + ) # get corner in global coordinates relative to box center corner = box_rot * corner - # compute distance to plane, skip if too far or pointing up + # compute distance to plane ldist = wp.dot(plane_normal, corner) - if center_dist + ldist > 0 or ldist > 0: - continue - cdist = center_dist + ldist - dist[ncontact] = cdist - pos[ncontact] = corner + box_pos - 0.5 * plane_normal * cdist - ncontact += 1 - - if ncontact >= 4: - break + dist[i] = cdist + pos[i] = corner + box_pos - 0.5 * plane_normal * cdist return dist, pos, plane_normal @@ -586,8 +580,7 @@ def box_box( box2_pos: Center position of the second box. box2_rot: Rotation matrix of the second box. box2_size: Half-extents of the second box along each axis. - margin: Distance threshold for early contact generation (default: 0.0). - When positive, contacts are generated before boxes overlap. + margin: Collision tolerance. Returns: - Vector of contact distances (wp.inf for unpopulated contacts). @@ -1111,10 +1104,8 @@ def capsule_box( halfaxis = axis * capsule_half_length # halfaxis is the capsule direction axisdir = wp.int32(halfaxis[0] > 0.0) + 2 * wp.int32(halfaxis[1] > 0.0) + 4 * wp.int32(halfaxis[2] > 0.0) - bestdistmax = 2.0 * (capsule_radius + capsule_half_length + box_size[0] + box_size[1] + box_size[2]) - # keep track of closest point - bestdist = wp.float32(bestdistmax) + bestdist = wp.float32(1.0e32) bestsegmentpos = wp.float32(-12) # cltype: encoded collision configuration @@ -1253,7 +1244,7 @@ def capsule_box( p = wp.vec2(pos.x, pos.y) dd = wp.vec2(halfaxis.x, halfaxis.y) - s = wp.vec2(box_size.x, box_size.y) + s = wp.vec2(box_size[0], box_size[1]) secondpos = wp.float32(-4.0) uu = dd.x * s.y 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 831565f9..210c7541 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py @@ -38,6 +38,7 @@ def _zero_constraint_counts( ne_weld_out: wp.array(dtype=int), ne_jnt_out: wp.array(dtype=int), ne_ten_out: wp.array(dtype=int), + ne_flex_out: wp.array(dtype=int), ): worldid = wp.tid() @@ -47,6 +48,7 @@ def _zero_constraint_counts( ne_weld_out[worldid] = 0 ne_jnt_out[worldid] = 0 ne_ten_out[worldid] = 0 + ne_flex_out[worldid] = 0 nf_out[worldid] = 0 nl_out[worldid] = 0 nefc_out[worldid] = 0 @@ -485,6 +487,81 @@ def _efc_equality_tendon( ) +@wp.kernel +def _efc_equality_flex( + # Model: + nv: int, + opt_timestep: wp.array(dtype=float), + flexedge_length0: wp.array(dtype=float), + flexedge_invweight0: wp.array(dtype=float), + eq_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_flex_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + flexedge_J_in: wp.array3d(dtype=float), + flexedge_length_in: wp.array2d(dtype=float), + njmax_in: int, + # In: + refsafe_in: int, + # Data out: + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ne_flex_out: wp.array(dtype=int), +): + worldid, eqflexid, edgeid = wp.tid() + eqid = eq_flex_adr[eqflexid] + + wp.atomic_add(ne_flex_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= njmax_in: + return + + pos = flexedge_length_in[worldid, edgeid] - flexedge_length0[edgeid] + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + + Jqvel = float(0.0) + for i in range(nv): + J = flexedge_J_in[worldid, edgeid, i] + efc_J_out[worldid, efcid, i] = J + Jqvel += J * qvel_in[worldid, i] + + _update_efc_row( + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + refsafe_in, + efcid, + pos, + pos, + flexedge_invweight0[edgeid], + solref, + solimp, + 0.0, + Jqvel, + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) + + @wp.kernel def _efc_friction_dof( # Model: @@ -1135,7 +1212,7 @@ def _efc_contact_pyramidal( # Model: nv: int, opt_timestep: wp.array(dtype=float), - opt_impratio: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), body_invweight0: wp.array2d(dtype=wp.vec2), @@ -1200,9 +1277,8 @@ def _efc_contact_pyramidal( contact_efc_address_out[conid, dimid] = -1 return - opt_timestep_id = worldid % opt_timestep.shape[0] - timestep = opt_timestep[opt_timestep_id] - impratio = opt_impratio[opt_timestep_id] + timestep = opt_timestep[worldid % opt_timestep.shape[0]] + impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] contact_efc_address_out[conid, dimid] = efcid geom = geom_in[conid] @@ -1223,7 +1299,7 @@ def _efc_contact_pyramidal( fri0 = friction[0] frii = friction[dimid2 - 1] invweight = invweight + fri0 * fri0 * invweight - invweight = invweight * 2.0 * fri0 * fri0 / impratio + invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt Jqvel = float(0.0) for i in range(nv): @@ -1306,7 +1382,7 @@ def _efc_contact_elliptic( # Model: nv: int, opt_timestep: wp.array(dtype=float), - opt_impratio: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), body_invweight0: wp.array2d(dtype=wp.vec2), @@ -1370,9 +1446,8 @@ def _efc_contact_elliptic( contact_efc_address_out[conid, dimid] = -1 return - opt_timestep_id = worldid % opt_timestep.shape[0] - timestep = opt_timestep[opt_timestep_id] - impratio = opt_impratio[opt_timestep_id] + timestep = opt_timestep[worldid % opt_timestep.shape[0]] + impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] contact_efc_address_out[conid, dimid] = efcid geom = geom_in[conid] @@ -1432,8 +1507,7 @@ def _efc_contact_elliptic( if solreffriction[0] or solreffriction[1]: ref = solreffriction - # TODO(team): precompute 1 / impratio - invweight = invweight / impratio + invweight = invweight * impratio_invsqrt * impratio_invsqrt friction = friction_in[conid] if dimid > 1: @@ -1482,11 +1556,12 @@ def _num_equality( ne_weld_in: wp.array(dtype=int), ne_jnt_in: wp.array(dtype=int), ne_ten_in: wp.array(dtype=int), + ne_flex_in: wp.array(dtype=int), # Data out: ne_out: wp.array(dtype=int), ): worldid = wp.tid() - ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne_flex_in[worldid] ne_out[worldid] = ne @@ -1496,7 +1571,7 @@ def make_constraint(m: types.Model, d: types.Data): wp.launch( _zero_constraint_counts, dim=d.nworld, - inputs=[d.ne, d.nf, d.nl, d.nefc, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten], + inputs=[d.ne, d.nf, d.nl, d.nefc, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex], ) if not (m.opt.disableflags & types.DisableBit.CONSTRAINT): @@ -1663,10 +1738,42 @@ def make_constraint(m: types.Model, d: types.Data): ], ) + wp.launch( + _efc_equality_flex, + dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge), + inputs=[ + m.nv, + m.opt.timestep, + m.flexedge_length0, + m.flexedge_invweight0, + m.eq_solref, + m.eq_solimp, + m.eq_flex_adr, + d.qvel, + d.flexedge_J, + d.flexedge_length, + d.njmax, + refsafe, + ], + outputs=[ + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + d.ne_flex, + ], + ) + wp.launch( _num_equality, dim=d.nworld, - inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten], + inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex], outputs=[d.ne], ) @@ -1847,7 +1954,7 @@ def make_constraint(m: types.Model, d: types.Data): inputs=[ m.nv, m.opt.timestep, - m.opt.impratio, + m.opt.impratio_invsqrt, m.body_parentid, m.body_rootid, m.body_invweight0, @@ -1892,7 +1999,7 @@ def make_constraint(m: types.Model, d: types.Data): inputs=[ m.nv, m.opt.timestep, - m.opt.impratio, + m.opt.impratio_invsqrt, m.body_parentid, m.body_rootid, m.body_invweight0, 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 83332751..84a2ce92 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py @@ -21,21 +21,18 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import DynType from mujoco.mjx.third_party.mujoco_warp._src.types import GainType from mujoco.mjx.third_party.mujoco_warp._src.types import Model +from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet 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 nested_kernel wp.set_module_options({"enable_backward": False}) -# TODO(team): improve performance with tile operations? @wp.kernel -def _qderiv_actuator_passive( +def _qderiv_actuator_passive_vel( # Model: - nu: int, - opt_timestep: wp.array(dtype=float), - opt_disableflags: int, - opt_is_sparse: bool, - dof_damping: wp.array2d(dtype=float), actuator_dyntype: wp.array(dtype=int), actuator_gaintype: wp.array(dtype=int), actuator_biastype: wp.array(dtype=int), @@ -46,9 +43,75 @@ def _qderiv_actuator_passive( # Data in: act_in: wp.array2d(dtype=float), ctrl_in: wp.array2d(dtype=float), + # Out: + vel_out: wp.array2d(dtype=float), +): + worldid, actid = wp.tid() + + actuator_gainprm_id = worldid % actuator_gainprm.shape[0] + actuator_biasprm_id = worldid % actuator_biasprm.shape[0] + + if actuator_gaintype[actid] == GainType.AFFINE: + gain = actuator_gainprm[actuator_gainprm_id, actid][2] + else: + gain = 0.0 + + if actuator_biastype[actid] == BiasType.AFFINE: + bias = actuator_biasprm[actuator_biasprm_id, actid][2] + else: + bias = 0.0 + + if bias == 0.0 and gain == 0.0: + vel_out[worldid, actid] = 0.0 + return + + vel = float(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] + + vel_out[worldid, actid] = vel + + +@cache_kernel +def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int): + @nested_kernel(module="unique", enable_backward=False) + def kernel( + # Data in: + vel_in: wp.array3d(dtype=float), + actuator_moment_in: wp.array3d(dtype=float), + # In: + adr: wp.array(dtype=int), + # Out: + qDeriv_out: wp.array3d(dtype=float), + ): + worldid, nodeid = wp.tid() + TILE_SIZE = wp.static(tile.size) + NU = wp.static(nu) + + dofid = adr[nodeid] + vel_tile = wp.tile_load(vel_in[worldid], shape=(NU, 1), bounds_check=False) + moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, TILE_SIZE), offset=(0, dofid), bounds_check=False) + moment_weighted = wp.tile_map(wp.mul, wp.tile_broadcast(vel_tile, shape=(NU, TILE_SIZE)), moment_tile) + qderiv_tile = wp.tile_matmul(wp.tile_transpose(moment_tile), moment_weighted) + wp.tile_store(qDeriv_out[worldid], qderiv_tile, offset=(dofid, dofid), bounds_check=False) + + return kernel + + +@wp.kernel +def _qderiv_actuator_passive_actuation_sparse( + # Model: + nu: int, + # Data in: actuator_moment_in: wp.array3d(dtype=float), - qM_in: wp.array3d(dtype=float), # In: + vel_in: wp.array2d(dtype=float), qMi: wp.array(dtype=int), qMj: wp.array(dtype=int), # Out: @@ -58,40 +121,45 @@ def _qderiv_actuator_passive( dofiid = qMi[elemid] dofjid = qMj[elemid] + qderiv_contrib = float(0.0) + for actid in range(nu): + vel = vel_in[worldid, actid] + if vel == 0.0: + continue - qderiv = float(0.0) - if not opt_disableflags & DisableBit.ACTUATION: - actuator_gainprm_id = worldid % actuator_gainprm.shape[0] - actuator_biasprm_id = worldid % actuator_biasprm.shape[0] + moment_i = actuator_moment_in[worldid, actid, dofiid] + moment_j = actuator_moment_in[worldid, actid, dofjid] - for actid in range(nu): - if actuator_gaintype[actid] == GainType.AFFINE: - gain = actuator_gainprm[actuator_gainprm_id, actid][2] - else: - gain = 0.0 + qderiv_contrib += moment_i * moment_j * vel - if actuator_biastype[actid] == BiasType.AFFINE: - bias = actuator_biasprm[actuator_biasprm_id, actid][2] - else: - bias = 0.0 + qDeriv_out[worldid, 0, elemid] = qderiv_contrib - if bias == 0.0 and gain == 0.0: - continue - 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] +@wp.kernel +def _qderiv_actuator_passive( + # Model: + opt_timestep: wp.array(dtype=float), + opt_disableflags: int, + opt_is_sparse: bool, + dof_damping: wp.array2d(dtype=float), + # Data in: + qM_in: wp.array3d(dtype=float), + # In: + qMi: wp.array(dtype=int), + qMj: wp.array(dtype=int), + qDeriv_in: wp.array3d(dtype=float), + # Out: + qDeriv_out: wp.array3d(dtype=float), +): + worldid, elemid = wp.tid() - if vel != 0.0: - qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel + dofiid = qMi[elemid] + dofjid = qMj[elemid] - # TODO(team): fluid model derivative + if opt_is_sparse: + qderiv = qDeriv_in[worldid, 0, elemid] + else: + qderiv = qDeriv_in[worldid, dofiid, dofjid] if not opt_disableflags & DisableBit.DAMPER and dofiid == dofjid: qderiv -= dof_damping[worldid % dof_damping.shape[0], dofiid] @@ -143,53 +211,82 @@ def _qderiv_tendon_damping( @event_scope -def deriv_smooth_vel(m: Model, d: Data, qDeriv: wp.array2d(dtype=float)): +def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)): """Analytical derivative of smooth forces w.r.t. velocities. Args: m: The model containing kinematic and dynamic information (device). d: The data object containing the current state and output arrays (device). - qDeriv: Analytical derivative of smooth forces w.r.t. velocity. + out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities). """ - qMi = m.qM_fullm_i if m.opt.is_sparse else m.dof_tri_row - qMj = m.qM_fullm_j if m.opt.is_sparse else m.dof_tri_col + qMi = m.qM_fullm_i + qMj = m.qM_fullm_j + + # TODO(team): implicit requires different sparsity structure if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)): + # TODO(team): only clear elements not set by _qderiv_actuator_passive + out.zero_() + if m.nu > 0 and not m.opt.disableflags & DisableBit.ACTUATION: + vel = wp.empty((d.nworld, m.nu), dtype=float) + wp.launch( + _qderiv_actuator_passive_vel, + dim=(d.nworld, m.nu), + inputs=[ + m.actuator_dyntype, + m.actuator_gaintype, + m.actuator_biastype, + m.actuator_actadr, + m.actuator_actnum, + m.actuator_gainprm, + m.actuator_biasprm, + d.act, + d.ctrl, + ], + outputs=[vel], + ) + if m.opt.is_sparse: + wp.launch( + _qderiv_actuator_passive_actuation_sparse, + dim=(d.nworld, qMi.size), + inputs=[m.nu, d.actuator_moment, vel, qMi, qMj], + outputs=[out], + ) + else: + vel_3d = vel.reshape(vel.shape + (1,)) + for tile in m.qM_tiles: + wp.launch_tiled( + _qderiv_actuator_passive_actuation_dense(tile, m.nu), + dim=(d.nworld, tile.adr.size), + inputs=[vel_3d, d.actuator_moment, tile.adr], + outputs=[out], + block_dim=m.block_dim.mul_m_dense, + ) wp.launch( _qderiv_actuator_passive, dim=(d.nworld, qMi.size), inputs=[ - m.nu, m.opt.timestep, m.opt.disableflags, m.opt.is_sparse, m.dof_damping, - m.actuator_dyntype, - m.actuator_gaintype, - m.actuator_biastype, - m.actuator_actadr, - m.actuator_actnum, - m.actuator_gainprm, - m.actuator_biasprm, - d.act, - d.ctrl, - d.actuator_moment, d.qM, qMi, qMj, + out, ], - outputs=[qDeriv], + outputs=[out], ) else: # TODO(team): directly utilize qM for these settings - wp.copy(qDeriv, d.qM) + wp.copy(out, d.qM) if not m.opt.disableflags & DisableBit.DAMPER: wp.launch( _qderiv_tendon_damping, dim=(d.nworld, qMi.size), inputs=[m.ntendon, m.opt.timestep, m.opt.is_sparse, m.tendon_damping, d.ten_J, qMi, qMj], - outputs=[qDeriv], + outputs=[out], ) # TODO(team): rne 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 41df144d..4f8a0501 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -509,6 +509,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True): smooth.kinematics(m, d) smooth.com_pos(m, d) smooth.camlight(m, d) + smooth.flex(m, d) smooth.tendon(m, d) smooth.crb(m, d) smooth.tendon_armature(m, d) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py index 00d58af7..00449b86 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py @@ -121,7 +121,9 @@ def inv_constraint(m: Model, d: Data): return # update - solver.create_context(m, d, grad=False) + h = wp.empty((d.nworld, 0, 0), dtype=float) # not used + hfactor = wp.empty((d.nworld, 0, 0), dtype=float) # not used + solver.create_context(m, d, h, hfactor, grad=False) def inverse(m: Model, d: Data): 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 2af891c5..faedac78 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -49,6 +49,16 @@ def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> Union[wp. return array +def is_sparse(mjm: mujoco.MjModel) -> bool: + if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO: + if mjm.nv > 32: + return True + else: + return False + else: + return bool(mujoco.mj_isSparse(mjm)) + + def put_model(mjm: mujoco.MjModel) -> types.Model: """Creates a model on device. @@ -93,9 +103,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: 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.") - if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any(): raise NotImplementedError("Flex collisions are not implemented.") @@ -137,23 +144,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: 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.") + 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, types.GeomType.BOX) and not_implemented(reftype, refid, types.GeomType.BOX): + raise NotImplementedError(f"Collision sensors with box-box collisions are not implemented.") # create opt - opt = types.Option(**{f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)}) + opt_kwargs = {f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)} + if hasattr(mjm.opt, "impratio"): + opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL)) + opt = types.Option(**opt_kwargs) # 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 @@ -161,7 +165,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: opt.tolerance = max(opt.tolerance, 1e-6) # warp only fields - opt.is_sparse = bool(mujoco.mj_isSparse(mjm)) + opt.is_sparse = is_sparse(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 @@ -170,7 +174,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: 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 + contact_sensor_maxmatch_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "contact_sensor_maxmatch") + if contact_sensor_maxmatch_id > -1: + opt.contact_sensor_maxmatch = mjm.numeric_data[mjm.numeric_adr[contact_sensor_maxmatch_id]] + else: + opt.contact_sensor_maxmatch = 64 # place opt on device for f in dataclasses.fields(types.Option): @@ -188,6 +196,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.opt = opt m.stat = stat + m.nv_pad = _get_padded_sizes( + mjm.nv, 0, is_sparse(mjm), types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE + )[1] 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() @@ -249,10 +260,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: nxn_pairid_contact[upper_tri_index(mjm.ngeom, mjm.pair_geom1[i], mjm.pair_geom2[i])] = i 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)) + nxn_pairid_collision = -1 * np.ones(len(geom1), dtype=int) pairids = [] - collision_geom_adr = [0] - m.sensor_collision_start_adr = [] + sensor_collision_start_adr = [] for i in range(sensor_collision_adr.size): sensorid = sensor_collision_adr[i] objtype = mjm.sensor_objtype[sensorid] @@ -275,24 +288,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: id2 = refid # collide all pairs - geomid = 0 for geom1id in range(id1, id1 + n1): for geom2id in range(id2, id2 + n2): pairid = upper_tri_index(mjm.ngeom, geom1id, geom2id) if pairid in pairids: - m.sensor_collision_start_adr.append(nxn_pairid_collision[pairid]) + sensor_collision_start_adr.append(nxn_pairid_collision[pairid]) else: + npairids = len(pairids) + nxn_pairid_collision[pairid] = npairids + sensor_collision_start_adr.append(npairids) pairids.append(pairid) - adr = collision_geom_adr[-1] + geomid - nxn_pairid_collision[pairid] = 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) m.nsensorcollision = (nxn_pairid_collision >= 0).sum() + m.sensor_collision_start_adr = np.array(sensor_collision_start_adr) nxn_include = (nxn_pairid_contact > -2) | (nxn_pairid_collision >= 0) if nxn_include.sum() < 250_000: @@ -363,6 +372,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: 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] + m.eq_flex_adr = np.nonzero(mjm.eq_type == types.EqType.FLEX)[0] # fixed tendon m.tendon_jnt_adr, m.wrap_jnt_adr = [], [] @@ -545,7 +555,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) - nv_padded = round_up(nv, tile_size) if is_sparse else round_up(nv, 4) + nv_padded = round_up(nv, tile_size) if (is_sparse or nv > 32) else round_up(nv, 4) return njmax_padded, nv_padded @@ -592,8 +602,8 @@ def make_data( 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) + tile_size = types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE + sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, is_sparse(mjm), tile_size) sizes["nworld"] = nworld sizes["naconmax"] = naconmax sizes["njmax"] = njmax @@ -601,7 +611,18 @@ def make_data( 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)}) + # world body and static geom (attached to the world) poses are precomputed + # 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) + + # mocap + mocap_body = np.nonzero(mjm.body_mocapid >= 0)[0] + mocap_id = mjm.body_mocapid[mocap_body] + d_kwargs = { + "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float), "contact": contact, "efc": efc, "nworld": nworld, @@ -609,8 +630,20 @@ def make_data( "njmax": njmax, "qM": None, "qLD": None, - "geom_xpos": None, - "geom_xmat": None, + # world body + "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat), + "xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33), + "ximat": wp.array(np.tile(mjd.ximat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33), + # static geoms + "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), + # mocap + "mocap_pos": wp.array(np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.vec3), + "mocap_quat": wp.array( + np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.quat + ), + # equality constraints + "eq_active": wp.array(np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), shape=(nworld, mjm.neq), dtype=bool), } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -619,21 +652,13 @@ def make_data( d = types.Data(**d_kwargs) - if mujoco.mj_isSparse(mjm): + if is_sparse(mjm): d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float) d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float) else: 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) - 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 d @@ -691,8 +716,8 @@ def put_data( 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) + tile_size = types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE + sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, is_sparse(mjm), tile_size) sizes["nworld"] = nworld sizes["naconmax"] = naconmax sizes["njmax"] = njmax @@ -766,11 +791,13 @@ def put_data( "qLD": None, "ten_J": None, "actuator_moment": None, + "flexedge_J": None, "nacon": None, "ne_connect": None, "ne_weld": None, "ne_jnt": None, "ne_ten": None, + "ne_flex": None, "nsolving": None, } for f in dataclasses.fields(types.Data): @@ -785,12 +812,9 @@ def put_data( d = types.Data(**d_kwargs) d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int) - if mujoco.mj_isSparse(mjm): + if is_sparse(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)) - 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) @@ -799,8 +823,21 @@ def put_data( 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) + + if mujoco.mj_isSparse(mjm): + 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), 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) + 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) # TODO(taylorhowell): sparse actuator_moment actuator_moment = np.zeros((mjm.nu, mjm.nv)) @@ -812,6 +849,7 @@ def put_data( 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.ne_flex = wp.full(nworld, np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX) & mjd.eq_active), dtype=int) d.nsolving = wp.array([nworld], dtype=int) return d @@ -913,6 +951,7 @@ 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] 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] @@ -957,12 +996,9 @@ def get_data_into( result.contact.geom[:ncon] = d.contact.geom.numpy()[ncon_filter] result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon] - if mujoco.mj_isSparse(mjm): + if is_sparse(mjm): result.qM[:] = d.qM.numpy()[world_id, 0] result.qLD[:] = d.qLD.numpy()[world_id, 0] - if nefc > 0: - efc_J = d.efc.J.numpy()[world_id, efc_idx, : mjm.nv] - mujoco.mju_dense2sparse(result.efc_J, efc_J, result.efc_J_rownnz, result.efc_J_rowadr, result.efc_J_colind) else: qM = d.qM.numpy()[world_id] adr = 0 @@ -973,7 +1009,12 @@ def get_data_into( j = mjm.dof_parentid[j] adr += 1 mujoco.mj_factorM(mjm, result) - if nefc > 0: + + if nefc > 0: + if mujoco.mj_isSparse(mjm): + efc_J = d.efc.J.numpy()[world_id, efc_idx, : mjm.nv] + mujoco.mju_dense2sparse(result.efc_J, efc_J, result.efc_J_rownnz, result.efc_J_rowadr, result.efc_J_colind) + else: result.efc_J[: nefc * mjm.nv] = d.efc.J.numpy()[world_id, :nefc, : mjm.nv].flatten() # efc @@ -1072,6 +1113,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): ne_weld_out: wp.array(dtype=int), ne_jnt_out: wp.array(dtype=int), ne_ten_out: wp.array(dtype=int), + ne_flex_out: wp.array(dtype=int), nsolving_out: wp.array(dtype=int), ): worldid = wp.tid() @@ -1088,6 +1130,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): ne_weld_out[worldid] = 0 ne_jnt_out[worldid] = 0 ne_ten_out[worldid] = 0 + ne_flex_out[worldid] = 0 nf_out[worldid] = 0 nl_out[worldid] = 0 nefc_out[worldid] = 0 @@ -1095,8 +1138,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): nsolving_out[0] = nworld_in time_out[worldid] = 0.0 energy_out[worldid] = wp.vec2(0.0, 0.0) + qpos0_id = worldid % qpos0.shape[0] for i in range(nq): - qpos_out[worldid, i] = qpos0[worldid, i] + qpos_out[worldid, i] = qpos0[qpos0_id, i] if i < nv: qvel_out[worldid, i] = 0.0 qacc_warmstart_out[worldid, i] = 0.0 @@ -1254,6 +1298,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): d.ne_weld, d.ne_jnt, d.ne_ten, + d.ne_flex, d.nsolving, ], ) 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 85a38afb..61127a8b 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -551,11 +551,14 @@ def _qfrc_passive( @wp.kernel def _flex_elasticity( # Model: + nflex: int, opt_timestep: wp.array(dtype=float), body_dofadr: wp.array(dtype=int), flex_dim: wp.array(dtype=int), flex_vertadr: wp.array(dtype=int), flex_edgeadr: wp.array(dtype=int), + flex_elemadr: wp.array(dtype=int), + flex_elemnum: wp.array(dtype=int), flex_elemedgeadr: wp.array(dtype=int), flex_vertbodyid: wp.array(dtype=int), flex_elem: wp.array(dtype=int), @@ -574,22 +577,27 @@ def _flex_elasticity( ): worldid, elemid = wp.tid() timestep = opt_timestep[worldid % opt_timestep.shape[0]] - f = 0 # TODO(quaglino): this should become a function of t + + for i in range(nflex): + locid = elemid - flex_elemadr[i] + if locid >= 0 and locid < flex_elemnum[i]: + f = i + break dim = flex_dim[f] nvert = dim + 1 nedge = nvert * (nvert - 1) / 2 edges = wp.where( dim == 3, - wp.mat(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int), - wp.mat(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int), + 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), ) if timestep > 0.0 and not dsbl_damper: kD = flex_damping[f] / timestep else: kD = 0.0 - gradient = wp.mat(0.0, shape=(6, 6)) + gradient = wp.types.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]] @@ -601,14 +609,14 @@ def _flex_elasticity( elongation = wp.spatial_vectorf(0.0) for e in range(nedge): - idx = flex_elemedge[flex_elemedgeadr[f] + elemid * nedge + e] + idx = flex_elemedge[elemid * nedge + e] vel = flexedge_velocity_in[worldid, flex_edgeadr[f] + idx] deformed = flexedge_length_in[worldid, flex_edgeadr[f] + idx] reference = flexedge_length0[flex_edgeadr[f] + idx] previous = deformed - vel * timestep elongation[e] = deformed * deformed - reference * reference + (deformed * deformed - previous * previous) * kD - metric = wp.mat(0.0, shape=(6, 6)) + metric = wp.types.matrix(0.0, shape=(6, 6)) id = int(0) for ed1 in range(nedge): for ed2 in range(ed1, nedge): @@ -616,7 +624,7 @@ def _flex_elasticity( metric[ed2, ed1] = flex_stiffness[elemid, id] id += 1 - force = wp.mat(0.0, shape=(6, 3)) + force = wp.types.matrix(0.0, shape=(6, 3)) for ed1 in range(nedge): for ed2 in range(nedge): for i in range(2): @@ -633,10 +641,12 @@ def _flex_elasticity( @wp.kernel def _flex_bending( # Model: + nflex: int, body_dofadr: wp.array(dtype=int), flex_dim: wp.array(dtype=int), flex_vertadr: wp.array(dtype=int), flex_edgeadr: wp.array(dtype=int), + flex_edgenum: wp.array(dtype=int), flex_vertbodyid: wp.array(dtype=int), flex_edge: wp.array(dtype=wp.vec2i), flex_edgeflap: wp.array(dtype=wp.vec2i), @@ -648,22 +658,27 @@ def _flex_bending( ): worldid, edgeid = wp.tid() nvert = 4 - f = 0 # TODO(quaglino): this should become a function of t + + for i in range(nflex): + locid = edgeid - flex_edgeadr[i] + if locid >= 0 and locid < flex_edgenum[i]: + f = i + break if flex_dim[f] != 2: return - v = wp.vec4i( - flex_edge[edgeid + flex_edgeadr[f]][0], - flex_edge[edgeid + flex_edgeadr[f]][1], - flex_edgeflap[edgeid + flex_edgeadr[f]][0], - flex_edgeflap[edgeid + flex_edgeadr[f]][1], - ) - - if v[3] == -1: + if flex_edgeflap[edgeid][1] == -1: return - frc = wp.mat(0.0, shape=(4, 3)) + v = wp.vec4i( + flex_vertadr[f] + flex_edge[edgeid][0], + flex_vertadr[f] + flex_edge[edgeid][1], + flex_vertadr[f] + flex_edgeflap[edgeid][0], + flex_vertadr[f] + flex_edgeflap[edgeid][1], + ) + + frc = wp.types.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]] @@ -674,7 +689,7 @@ def _flex_bending( frc[3] = wp.cross(v1 - v0, v2 - v0) frc[0] = -(frc[1] + frc[2] + frc[3]) - force = wp.mat(0.0, shape=(nvert, 3)) + force = wp.types.matrix(0.0, shape=(nvert, 3)) for i in range(nvert): for x in range(3): for j in range(nvert): @@ -743,11 +758,14 @@ def passive(m: Model, d: Data): _flex_elasticity, dim=(d.nworld, m.nflexelem), inputs=[ + m.nflex, m.opt.timestep, m.body_dofadr, m.flex_dim, m.flex_vertadr, m.flex_edgeadr, + m.flex_elemadr, + m.flex_elemnum, m.flex_elemedgeadr, m.flex_vertbodyid, m.flex_elem, @@ -766,10 +784,12 @@ def passive(m: Model, d: Data): _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, 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 54f1bb48..36c50c6e 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -455,6 +455,7 @@ def _clock(time_in: wp.array(dtype=float), worldid: int) -> float: @wp.kernel def _sensor_pos( # Model: + ngeom: int, opt_magnetic: wp.array(dtype=wp.vec3), body_geomnum: wp.array(dtype=int), body_geomadr: wp.array(dtype=int), @@ -481,10 +482,9 @@ def _sensor_pos( sensor_refid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), + nxn_pairid: wp.array(dtype=wp.vec2i), sensor_pos_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), energy_in: wp.array(dtype=wp.vec2), @@ -605,13 +605,10 @@ def _sensor_pos( refid = sensor_refid[sensorid] # initialize - dist = float(1.0e32) + dist = float(sensor_cutoff[sensorid]) pnts = vec6(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) flip = bool(False) - collision_sensorid = collision_sensor_adr[sensorid] - collision_start_adr = sensor_collision_start_adr[collision_sensorid] - # check for flip direction if objtype == int(ObjType.BODY.value): n1 = body_geomnum[objid] @@ -627,12 +624,20 @@ def _sensor_pos( id2 = refid for geom1 in range(n1): + geomid1 = id1 + geom1 for geom2 in range(n2): - collisionid = collision_start_adr + geom1 * n2 + geom2 + geomid2 = id2 + geom2 + + if geomid1 <= geomid2: + pairid = math.upper_tri_index(ngeom, geomid1, geomid2) + else: + pairid = math.upper_tri_index(ngeom, geomid2, geomid1) + collisionid = nxn_pairid[pairid][1] + for i in range(8): dist_new = sensor_collision_in[worldid, collisionid, i, 0] - if dist_new <= dist: + if dist_new < dist: dist = dist_new if sensortype == SensorType.GEOMNORMAL or sensortype == SensorType.GEOMFROMTO: @@ -645,14 +650,12 @@ def _sensor_pos( sensor_collision_in[worldid, collisionid, i, 6], ) - geomid1 = id1 + geom1 - geomid2 = id2 + geom2 - - if geom_type[geomid2] < geom_type[geomid1]: - flip = True - elif geom_type[geomid1] == geom_type[geomid2]: - if geomid2 < geomid1: - flip = True + if geom_type[geomid1] > geom_type[geomid2]: + flip = True + elif geom_type[geomid1] == geom_type[geomid2]: + flip = geomid1 > geomid2 + else: + flip = False if sensortype == int(SensorType.GEOMDIST.value): _write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, dist, out) elif sensortype == int(SensorType.GEOMNORMAL.value): @@ -704,6 +707,7 @@ def _sensor_pos( def _sensor_collision( # Model: ngeom: int, + nxn_pairid: wp.array(dtype=wp.vec2i), # Data in: contact_dist_in: wp.array(dtype=float), contact_pos_in: wp.array(dtype=wp.vec3), @@ -713,7 +717,6 @@ def _sensor_collision( contact_type_in: wp.array(dtype=int), contact_geomcollisionid_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int), - collision_pairid_in: wp.array(dtype=wp.vec2i), # Out: sensor_collision_out: wp.array4d(dtype=float), ): @@ -732,7 +735,7 @@ def _sensor_collision( pairid = math.upper_tri_index(ngeom, geom[1], geom[0]) worldid = contact_worldid_in[conid] - collisionid = collision_pairid_in[pairid][1] + collisionid = nxn_pairid[pairid][1] geomcollisionid = contact_geomcollisionid_in[conid] dist = contact_dist_in[conid] @@ -800,6 +803,7 @@ def sensor_pos(m: Model, d: Data): dim=d.naconmax, inputs=[ m.ngeom, + m.nxn_pairid, d.contact.dist, d.contact.pos, d.contact.frame, @@ -808,7 +812,6 @@ def sensor_pos(m: Model, d: Data): d.contact.type, d.contact.geomcollisionid, d.nacon, - d.collision_pairid, ], outputs=[sensor_collision], ) @@ -817,6 +820,7 @@ def sensor_pos(m: Model, d: Data): _sensor_pos, dim=(d.nworld, m.sensor_pos_adr.size), inputs=[ + m.ngeom, m.opt.magnetic, m.body_geomnum, m.body_geomadr, @@ -843,10 +847,9 @@ def sensor_pos(m: Model, d: Data): m.sensor_refid, m.sensor_adr, m.sensor_cutoff, + m.nxn_pairid, m.sensor_pos_adr, m.rangefinder_sensor_adr, - m.sensor_collision_start_adr, - m.collision_sensor_adr, d.time, d.energy, d.qpos, 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 9e56f1f7..2bc68aff 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py @@ -40,28 +40,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) -@wp.kernel -def _kinematics_root( - # Data out: - xpos_out: wp.array2d(dtype=wp.vec3), - xquat_out: wp.array2d(dtype=wp.quat), - xmat_out: wp.array2d(dtype=wp.mat33), - xipos_out: wp.array2d(dtype=wp.vec3), - ximat_out: wp.array2d(dtype=wp.mat33), -): - worldid = wp.tid() - xpos_out[worldid, 0] = wp.vec3(0.0) - xquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) - xipos_out[worldid, 0] = wp.vec3(0.0) - xmat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32) - ximat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32) - - @wp.kernel def _kinematics_level( # Model: qpos0: wp.array2d(dtype=float), body_parentid: wp.array(dtype=int), + body_mocapid: wp.array(dtype=int), body_jntnum: wp.array(dtype=int), body_jntadr: wp.array(dtype=int), body_pos: wp.array2d(dtype=wp.vec3), @@ -74,6 +58,8 @@ def _kinematics_level( jnt_axis: wp.array2d(dtype=wp.vec3), # Data in: qpos_in: wp.array2d(dtype=float), + mocap_pos_in: wp.array2d(dtype=wp.vec3), + mocap_quat_in: wp.array2d(dtype=wp.quat), xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xmat_in: wp.array2d(dtype=wp.mat33), @@ -97,12 +83,12 @@ def _kinematics_level( body_quat_id = worldid % body_quat.shape[0] jnt_axis_id = worldid % jnt_axis.shape[0] - if jntnum == 0: - # no joints - apply fixed translation and rotation relative to parent - pid = body_parentid[bodyid] - xpos = (xmat_in[worldid, pid] * body_pos[body_pos_id, bodyid]) + xpos_in[worldid, pid] - xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[body_quat_id, bodyid]) - elif jntnum == 1 and jnt_type[jntadr] == JointType.FREE: + free_joint = False + if jntnum == 1: + jnt_type_ = jnt_type[jntadr] + free_joint = jnt_type_ == JointType.FREE + + if free_joint: # free joint qadr = jnt_qposadr[jntadr] xpos = wp.vec3(qpos[qadr], qpos[qadr + 1], qpos[qadr + 2]) @@ -116,8 +102,19 @@ def _kinematics_level( qpos0_id = worldid % qpos0.shape[0] jnt_pos_id = worldid % jnt_pos.shape[0] pid = body_parentid[bodyid] - xpos = (xmat_in[worldid, pid] * body_pos[body_pos_id, bodyid]) + xpos_in[worldid, pid] - xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[body_quat_id, bodyid]) + + # mocap bodies have world body as parent + mocapid = body_mocapid[bodyid] + if mocapid >= 0: + xpos = mocap_pos_in[worldid, mocapid] + xquat = mocap_quat_in[worldid, mocapid] + else: + xpos = body_pos[body_pos_id, bodyid] + xquat = body_quat[body_quat_id, bodyid] + + if pid >= 0: + xpos = xmat_in[worldid, pid] @ xpos + xpos_in[worldid, pid] + xquat = math.mul_quat(xquat_in[worldid, pid], xquat) for _ in range(jntnum): qadr = jnt_qposadr[jntadr] @@ -146,7 +143,8 @@ def _kinematics_level( jntadr += 1 xpos_out[worldid, bodyid] = xpos - xquat_out[worldid, bodyid] = wp.normalize(xquat) + xquat = wp.normalize(xquat) + xquat_out[worldid, bodyid] = xquat xmat_out[worldid, bodyid] = math.quat_to_mat(xquat) xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], xquat) ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid])) @@ -219,19 +217,33 @@ def _flex_vertices( @wp.kernel def _flex_edges( # Model: + nv: int, + nflex: int, + body_parentid: wp.array(dtype=int), + body_rootid: wp.array(dtype=int), body_dofadr: wp.array(dtype=int), + dof_bodyid: wp.array(dtype=int), flex_vertadr: wp.array(dtype=int), + flex_edgeadr: wp.array(dtype=int), + flex_edgenum: wp.array(dtype=int), flex_vertbodyid: wp.array(dtype=int), flex_edge: wp.array(dtype=wp.vec2i), # Data in: qvel_in: wp.array2d(dtype=float), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cdof_in: wp.array2d(dtype=wp.spatial_vector), flexvert_xpos_in: wp.array2d(dtype=wp.vec3), # Data out: + flexedge_J_out: wp.array3d(dtype=float), flexedge_length_out: wp.array2d(dtype=float), flexedge_velocity_out: wp.array2d(dtype=float), ): worldid, edgeid = wp.tid() - f = 0 # TODO(quaglino): get f from edgeid + for i in range(nflex): + locid = edgeid - flex_edgeadr[i] + if locid >= 0 and locid < flex_edgenum[i]: + f = i + break vbase = flex_vertadr[f] v = flex_edge[edgeid] pos1 = flexvert_xpos_in[worldid, vbase + v[0]] @@ -240,38 +252,20 @@ def _flex_edges( vecnorm = wp.length(vec) flexedge_length_out[worldid, edgeid] = vecnorm # TODO(quaglino): use Jacobian - i = body_dofadr[flex_vertbodyid[vbase + v[0]]] - j = body_dofadr[flex_vertbodyid[vbase + v[1]]] + b1 = flex_vertbodyid[vbase + v[0]] + b2 = flex_vertbodyid[vbase + v[1]] + i = body_dofadr[b1] + j = body_dofadr[b2] vel1 = wp.vec3(qvel_in[worldid, i], qvel_in[worldid, i + 1], qvel_in[worldid, i + 2]) vel2 = wp.vec3(qvel_in[worldid, j], qvel_in[worldid, j + 1], qvel_in[worldid, j + 2]) - flexedge_velocity_out[worldid, edgeid] = math.safe_div(wp.dot(vel2 - vel1, vec), vecnorm) - - -@wp.kernel -def _mocap( - # Model: - body_ipos: wp.array2d(dtype=wp.vec3), - body_iquat: wp.array2d(dtype=wp.quat), - mocap_bodyid: wp.array(dtype=int), - # Data in: - mocap_pos_in: wp.array2d(dtype=wp.vec3), - mocap_quat_in: wp.array2d(dtype=wp.quat), - # Data out: - xpos_out: wp.array2d(dtype=wp.vec3), - xquat_out: wp.array2d(dtype=wp.quat), - xmat_out: wp.array2d(dtype=wp.mat33), - xipos_out: wp.array2d(dtype=wp.vec3), - ximat_out: wp.array2d(dtype=wp.mat33), -): - worldid, mocapid = wp.tid() - bodyid = mocap_bodyid[mocapid] - mocap_quat = wp.normalize(mocap_quat_in[worldid, mocapid]) - xpos = mocap_pos_in[worldid, mocapid] - xpos_out[worldid, bodyid] = xpos - xquat_out[worldid, bodyid] = mocap_quat - xmat_out[worldid, bodyid] = math.quat_to_mat(mocap_quat) - xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], mocap_quat) - ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(mocap_quat, body_iquat[worldid % body_iquat.shape[0], bodyid])) + edge = wp.normalize(vec) + flexedge_velocity_out[worldid, edgeid] = wp.dot(vel2 - vel1, edge) + # Edge jacobian + for k in range(nv): + jacp1, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, k, worldid) + jacp2, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, k, worldid) + jacdif = jacp2 - jacp1 + flexedge_J_out[worldid, edgeid, k] = wp.dot(jacdif, edge) @event_scope @@ -282,8 +276,6 @@ def kinematics(m: Model, d: Data): derived positions and orientations of geoms, sites, and flexible elements, based on the current joint positions and any attached mocap bodies. """ - wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat]) - for i in range(1, len(m.body_tree)): body_tree = m.body_tree[i] wp.launch( @@ -292,6 +284,7 @@ def kinematics(m: Model, d: Data): inputs=[ m.qpos0, m.body_parentid, + m.body_mocapid, m.body_jntnum, m.body_jntadr, m.body_pos, @@ -303,6 +296,8 @@ def kinematics(m: Model, d: Data): m.jnt_pos, m.jnt_axis, d.qpos, + d.mocap_pos, + d.mocap_quat, d.xpos, d.xquat, d.xmat, @@ -311,13 +306,6 @@ def kinematics(m: Model, d: Data): outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat, d.xanchor, d.xaxis], ) - wp.launch( - _mocap, - dim=(d.nworld, m.nmocap), - inputs=[m.body_ipos, m.body_iquat, m.mocap_bodyid, d.mocap_pos, d.mocap_quat], - outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat], - ) - wp.launch( _geom_local_to_global, dim=(d.nworld, m.ngeom), @@ -332,12 +320,31 @@ def kinematics(m: Model, d: Data): outputs=[d.site_xpos, d.site_xmat], ) + +@event_scope +def flex(m: Model, d: Data): wp.launch(_flex_vertices, dim=(d.nworld, m.nflexvert), inputs=[m.flex_vertbodyid, d.xpos], outputs=[d.flexvert_xpos]) wp.launch( _flex_edges, dim=(d.nworld, m.nflexedge), - inputs=[m.body_dofadr, m.flex_vertadr, m.flex_vertbodyid, m.flex_edge, d.qvel, d.flexvert_xpos], - outputs=[d.flexedge_length, d.flexedge_velocity], + inputs=[ + m.nv, + m.nflex, + m.body_parentid, + m.body_rootid, + m.body_dofadr, + m.dof_bodyid, + m.flex_vertadr, + m.flex_edgeadr, + m.flex_edgenum, + m.flex_vertbodyid, + m.flex_edge, + d.qvel, + d.subtree_com, + d.cdof, + d.flexvert_xpos, + ], + outputs=[d.flexedge_J, d.flexedge_length, d.flexedge_velocity], ) 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 29e7f4ba..d8853115 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -31,10 +31,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) +_BLOCK_CHOLESKY_DIM = 32 + @wp.func def _rescale(nv: int, stat_meaninertia: float, value: float) -> float: - return value / (stat_meaninertia * float(wp.max(1, nv))) + return value / (stat_meaninertia * float(nv)) @wp.func @@ -291,10 +293,10 @@ def _eval( def linesearch_iterative( # Model: nv: int, - opt_impratio: wp.array(dtype=float), opt_tolerance: wp.array(dtype=float), opt_ls_tolerance: wp.array(dtype=float), opt_ls_iterations: int, + opt_impratio_invsqrt: wp.array(dtype=float), stat_meaninertia: float, # Data in: ne_in: wp.array(dtype=int), @@ -321,7 +323,7 @@ def linesearch_iterative( if efc_done_in[worldid]: return - impratio = opt_impratio[worldid % opt_impratio.shape[0]] + impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] efc_type = efc_type_in[worldid] efc_id = efc_id_in[worldid] efc_D = efc_D_in[worldid] @@ -335,11 +337,10 @@ def linesearch_iterative( ne_clip = min(njmax_in, ne_in[worldid]) nef_clip = min(njmax_in, ne_clip + nf_in[worldid]) nefc_clip = min(njmax_in, nefc_in[worldid]) - impratio_invsqrt = 1.0 / wp.sqrt(impratio) # Calculate p0 - snorm = wp.math.sqrt(efc_search_dot_in[worldid]) - scale = stat_meaninertia * wp.float(wp.max(1, nv)) + snorm = wp.sqrt(efc_search_dot_in[worldid]) + scale = stat_meaninertia * wp.float(nv) gtol = tolerance * ls_tolerance * snorm * scale p0 = wp.vec3(efc_quad_gauss[0], efc_quad_gauss[1], 2.0 * efc_quad_gauss[2]) p0 += _eval_init( @@ -461,10 +462,10 @@ def _linesearch_iterative(m: types.Model, d: types.Data): dim=d.nworld, inputs=[ m.nv, - m.opt.impratio, m.opt.tolerance, m.opt.ls_tolerance, m.opt.ls_iterations, + m.opt.impratio_invsqrt, m.stat.meaninertia, d.ne, d.nf, @@ -484,6 +485,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data): d.njmax, ], outputs=[d.efc.alpha], + block_dim=m.block_dim.linesearch_iterative, ) @@ -496,8 +498,8 @@ def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> f @wp.kernel def linesearch_parallel_fused( # Model: - opt_impratio: wp.array(dtype=float), opt_ls_iterations: int, + opt_impratio_invsqrt: wp.array(dtype=float), opt_ls_parallel_min_step: float, # Data in: ne_in: wp.array(dtype=int), @@ -569,7 +571,7 @@ def linesearch_parallel_fused( continue friction = contact_friction_in[conid] - mu = friction[0] / wp.sqrt(opt_impratio[worldid % opt_impratio.shape[0]]) + mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] # unpack quad efcid1 = contact_efc_address_in[conid, 1] @@ -649,8 +651,8 @@ def _linesearch_parallel(m: types.Model, d: types.Data, cost: wp.array2d(dtype=f linesearch_parallel_fused, dim=(d.nworld, m.opt.ls_iterations), inputs=[ - m.opt.impratio, m.opt.ls_iterations, + m.opt.impratio_invsqrt, m.opt.ls_parallel_min_step, d.ne, d.nf, @@ -768,7 +770,7 @@ def linesearch_prepare_gauss( @wp.kernel def linesearch_prepare_quad( # Model: - opt_impratio: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), # Data in: nefc_in: wp.array(dtype=int), contact_friction_in: wp.array(dtype=types.vec5), @@ -814,7 +816,7 @@ def linesearch_prepare_quad( dim = contact_dim_in[conid] friction = contact_friction_in[conid] - mu = friction[0] / wp.sqrt(opt_impratio[worldid]) + mu = friction[0] * opt_impratio_invsqrt[worldid] u0 = Jaref * mu v0 = jv * mu @@ -949,7 +951,7 @@ def _linesearch(m: types.Model, d: types.Data, cost: wp.array2d(dtype=float)): linesearch_prepare_quad, dim=(d.nworld, d.njmax), inputs=[ - m.opt.impratio, + m.opt.impratio_invsqrt, d.nefc, d.contact.friction, d.contact.dim, @@ -1061,7 +1063,7 @@ def update_constraint_init_cost( @wp.kernel def update_constraint_efc( # Model: - opt_impratio: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), # Data in: ne_in: wp.array(dtype=int), nf_in: wp.array(dtype=int), @@ -1133,7 +1135,7 @@ def update_constraint_efc( dim = contact_dim_in[conid] friction = contact_friction_in[conid] - mu = friction[0] / wp.sqrt(opt_impratio[worldid]) + mu = friction[0] * opt_impratio_invsqrt[worldid] efcid0 = contact_efc_address_in[conid, 0] if efcid0 < 0: @@ -1261,7 +1263,7 @@ def _update_constraint(m: types.Model, d: types.Data): update_constraint_efc, dim=(d.nworld, d.njmax), inputs=[ - m.opt.impratio, + m.opt.impratio_invsqrt, d.ne, d.nf, d.nefc, @@ -1349,8 +1351,8 @@ def update_gradient_set_h_qM_lower_sparse( # Data in: qM_in: wp.array3d(dtype=float), efc_done_in: wp.array(dtype=bool), - # Data out: - efc_h_out: wp.array3d(dtype=float), + # Out: + h_out: wp.array3d(dtype=float), ): worldid, elementid = wp.tid() @@ -1359,7 +1361,7 @@ def update_gradient_set_h_qM_lower_sparse( i = qM_fullm_i[elementid] j = qM_fullm_j[elementid] - efc_h_out[worldid, i, j] += qM_in[worldid, 0, elementid] + h_out[worldid, i, j] += qM_in[worldid, 0, elementid] @wp.func @@ -1390,8 +1392,8 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int): efc_D_in: wp.array2d(dtype=float), efc_state_in: wp.array2d(dtype=int), efc_done_in: wp.array(dtype=bool), - # Data out: - efc_h_out: wp.array3d(dtype=float), + # Out: + h_out: wp.array3d(dtype=float), ): worldid, elementid = wp.tid() @@ -1442,7 +1444,7 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int): # AD: setting bounds_check to True explicitly here because for some reason it was # slower to disable it. - wp.tile_store(efc_h_out[worldid], sum_val, offset=(offset_i, offset_j), bounds_check=True) + wp.tile_store(h_out[worldid], sum_val, offset=(offset_i, offset_j), bounds_check=True) return kernel @@ -1463,8 +1465,8 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int efc_D_in: wp.array2d(dtype=float), efc_state_in: wp.array2d(dtype=int), efc_done_in: wp.array(dtype=bool), - # Data out: - efc_h_out: wp.array3d(dtype=float), + # Out: + h_out: wp.array3d(dtype=float), ): worldid = wp.tid() @@ -1503,7 +1505,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int sum_val += wp.tile_matmul(J_ki, J_kj) - wp.tile_store(efc_h_out[worldid], sum_val, bounds_check=False) + wp.tile_store(h_out[worldid], sum_val, bounds_check=False) return kernel @@ -1512,7 +1514,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int @wp.kernel def update_gradient_JTCJ( # Model: - opt_impratio: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), dof_tri_row: wp.array(dtype=int), dof_tri_col: wp.array(dtype=int), # Data in: @@ -1532,8 +1534,8 @@ def update_gradient_JTCJ( # In: nblocks_perblock: int, dim_block: int, - # Data out: - efc_h_out: wp.array3d(dtype=float), + # Out: + h_out: wp.array3d(dtype=float), ): conid_start, elementid = wp.tid() @@ -1564,7 +1566,7 @@ def update_gradient_JTCJ( continue fri = contact_friction_in[conid] - mu = math.safe_div(fri[0], wp.sqrt(opt_impratio[worldid])) + mu = fri[0] * opt_impratio_invsqrt[worldid] mu2 = mu * mu dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2)) @@ -1589,7 +1591,7 @@ def update_gradient_JTCJ( t = wp.max(t, types.MJ_MINVAL) ttt = wp.max(t * t * t, types.MJ_MINVAL) - efc_h = float(0.0) + h = float(0.0) for dim1id in range(condim): if dim1id == 0: @@ -1641,12 +1643,12 @@ def update_gradient_JTCJ( hcone *= dm * fri1 * fri2 if hcone != 0.0: - efc_h += hcone * efc_J11 * efc_J22 + h += hcone * efc_J11 * efc_J22 if dim1id != dim2id: - efc_h += hcone * efc_J12 * efc_J21 + h += hcone * efc_J12 * efc_J21 - efc_h_out[worldid, dof1id, dof2id] += efc_h + h_out[worldid, dof1id, dof2id] += h @cache_kernel @@ -1655,7 +1657,7 @@ def update_gradient_cholesky(tile_size: int): def kernel( # Data in: efc_grad_in: wp.array2d(dtype=float), - efc_h_in: wp.array3d(dtype=float), + h_in: wp.array3d(dtype=float), efc_done_in: wp.array(dtype=bool), # Data out: efc_Mgrad_out: wp.array2d(dtype=float), @@ -1666,7 +1668,7 @@ def update_gradient_cholesky(tile_size: int): if efc_done_in[worldid]: return - mat_tile = wp.tile_load(efc_h_in[worldid], shape=(TILE_SIZE, TILE_SIZE)) + mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE)) fact_tile = wp.tile_cholesky(mat_tile) input_tile = wp.tile_load(efc_grad_in[worldid], shape=TILE_SIZE) output_tile = wp.tile_cholesky_solve(fact_tile, input_tile) @@ -1676,34 +1678,48 @@ def update_gradient_cholesky(tile_size: int): @cache_kernel -def update_gradient_cholesky_blocked(tile_size: int): +def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int): @nested_kernel(module="unique", enable_backward=False) def kernel( # Data in: efc_grad_in: wp.array3d(dtype=float), - efc_h_in: wp.array3d(dtype=float), + h_in: wp.array3d(dtype=float), efc_done_in: wp.array(dtype=bool), - matrix_size: int, - cholesky_L_tmp: wp.array3d(dtype=float), - cholesky_y_tmp: wp.array3d(dtype=float), + hfactor: wp.array3d(dtype=float), # Data out: efc_Mgrad_out: wp.array3d(dtype=float), ): - worldid, tid_block = wp.tid() + worldid = wp.tid() TILE_SIZE = wp.static(tile_size) if efc_done_in[worldid]: return - wp.static(create_blocked_cholesky_func(TILE_SIZE))(tid_block, efc_h_in[worldid], matrix_size, cholesky_L_tmp[worldid]) - wp.static(create_blocked_cholesky_solve_func(TILE_SIZE))( - tid_block, cholesky_L_tmp[worldid], efc_grad_in[worldid], cholesky_y_tmp[worldid], matrix_size, efc_Mgrad_out[worldid] + # We need matrix size both as a runtime input as well as a static input: + # static input is needed to specify the tile sizes for the compiler + # runtime input is needed for the loop bounds, otherwise warp will unroll + # unconditionally leading to shared memory capacity issues. + + wp.static(create_blocked_cholesky_func(TILE_SIZE))(h_in[worldid], matrix_size, hfactor[worldid]) + wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))( + hfactor[worldid], efc_grad_in[worldid], matrix_size, efc_Mgrad_out[worldid] ) return kernel -def _update_gradient(m: types.Model, d: types.Data): +@wp.kernel +def padding_h(nv: int, efc_done_in: wp.array(dtype=bool), h_out: wp.array3d(dtype=float)): + worldid, elementid = wp.tid() + + if efc_done_in[worldid]: + return + + dofid = nv + elementid + h_out[worldid, dofid, dofid] = 1.0 + + +def _update_gradient(m: types.Model, d: types.Data, h: wp.array3d(dtype=float), hfactor: wp.array3d(dtype=float)): # grad = Ma - qfrc_smooth - qfrc_constraint wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[d.efc.done], outputs=[d.efc.grad_dot]) @@ -1731,7 +1747,7 @@ def _update_gradient(m: types.Model, d: types.Data): d.efc.state, d.efc.done, ], - outputs=[d.efc.h], + outputs=[h], block_dim=m.block_dim.update_gradient_JTDAJ_sparse, ) @@ -1739,7 +1755,7 @@ def _update_gradient(m: types.Model, d: types.Data): update_gradient_set_h_qM_lower_sparse, dim=(d.nworld, m.qM_fullm_i.size), inputs=[m.qM_fullm_i, m.qM_fullm_j, d.qM, d.efc.done], - outputs=[d.efc.h], + outputs=[h], ) else: nv_padded = d.efc.J.shape[2] @@ -1754,7 +1770,7 @@ def _update_gradient(m: types.Model, d: types.Data): d.efc.state, d.efc.done, ], - outputs=[d.efc.h], + outputs=[h], block_dim=m.block_dim.update_gradient_JTDAJ_dense, ) @@ -1786,7 +1802,7 @@ def _update_gradient(m: types.Model, d: types.Data): update_gradient_JTCJ, dim=(dim_block, m.dof_tri_row.size), inputs=[ - m.opt.impratio, + m.opt.impratio_invsqrt, m.dof_tri_row, m.dof_tri_col, d.contact.dist, @@ -1805,32 +1821,32 @@ def _update_gradient(m: types.Model, d: types.Data): nblocks_perblock, dim_block, ], - outputs=[d.efc.h], + outputs=[h], ) # TODO(team): Define good threshold for blocked vs non-blocked cholesky - if m.nv < 32: + if m.nv <= _BLOCK_CHOLESKY_DIM: wp.launch_tiled( update_gradient_cholesky(m.nv), dim=d.nworld, - inputs=[d.efc.grad, d.efc.h, d.efc.done], + inputs=[d.efc.grad, h, d.efc.done], outputs=[d.efc.Mgrad], block_dim=m.block_dim.update_gradient_cholesky, ) else: + wp.launch( + padding_h, + dim=(d.nworld, m.nv_pad - m.nv), + inputs=[m.nv, d.efc.done], + outputs=[h], + ) + wp.launch_tiled( - update_gradient_cholesky_blocked(16), + update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad), dim=d.nworld, - inputs=[ - d.efc.grad.reshape(shape=(d.nworld, m.nv, 1)), - d.efc.h, - d.efc.done, - m.nv, - d.efc.cholesky_L_tmp, - d.efc.cholesky_y_tmp.reshape(shape=(d.nworld, m.nv, 1)), - ], - outputs=[d.efc.Mgrad.reshape(shape=(d.nworld, m.nv, 1))], - block_dim=m.block_dim.update_gradient_cholesky, + inputs=[d.efc.grad.reshape(shape=(d.nworld, d.efc.grad.shape[1], 1)), h, d.efc.done, hfactor], + outputs=[d.efc.Mgrad.reshape(shape=(d.nworld, d.efc.Mgrad.shape[1], 1))], + block_dim=m.block_dim.update_gradient_cholesky_blocked, ) else: raise ValueError(f"Unknown solver type: {m.opt.solver}") @@ -1951,7 +1967,7 @@ def solve_done( tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]] improvement = _rescale(nv, stat_meaninertia, efc_prev_cost_in[worldid] - efc_cost_in[worldid]) - gradient = _rescale(nv, stat_meaninertia, wp.math.sqrt(efc_grad_dot_in[worldid])) + gradient = _rescale(nv, stat_meaninertia, wp.sqrt(efc_grad_dot_in[worldid])) done = (improvement < tolerance) or (gradient < tolerance) if done or solver_niter_out[worldid] == opt_iterations: # if the solver has converged or the maximum number of iterations has been reached then @@ -1964,6 +1980,8 @@ def solve_done( def _solver_iteration( m: types.Model, d: types.Data, + h: wp.array3d(dtype=float), + hfactor: wp.array3d(dtype=float), step_size_cost: wp.array2d(dtype=float), ): _linesearch(m, d, step_size_cost) @@ -1977,7 +1995,7 @@ def _solver_iteration( ) _update_constraint(m, d) - _update_gradient(m, d) + _update_gradient(m, d, h, hfactor) # polak-ribiere if m.opt.solver == types.SolverType.CG: @@ -2014,7 +2032,9 @@ def _solver_iteration( ) -def create_context(m: types.Model, d: types.Data, grad: bool = True): +def create_context( + m: types.Model, d: types.Data, h: wp.array3d(dtype=float), hfactor: wp.array3d(dtype=float), grad: bool = True +): # initialize some efc arrays wp.launch( solve_init_efc, @@ -2036,7 +2056,7 @@ def create_context(m: types.Model, d: types.Data, grad: bool = True): _update_constraint(m, d) if grad: - _update_gradient(m, d) + _update_gradient(m, d, h, hfactor) @event_scope @@ -2055,8 +2075,19 @@ def _solve(m: types.Model, d: types.Data): else: wp.copy(d.qacc, d.qacc_smooth) + # Newton solver Hessian + if m.opt.solver == types.SolverType.NEWTON: + h = wp.zeros((d.nworld, m.nv_pad, m.nv_pad), dtype=float) + if m.nv > _BLOCK_CHOLESKY_DIM: + hfactor = wp.zeros((d.nworld, m.nv_pad, m.nv_pad), dtype=float) + else: + hfactor = wp.empty((d.nworld, 0, 0), dtype=float) + else: + h = wp.empty((d.nworld, 0, 0), dtype=float) + hfactor = wp.empty((d.nworld, 0, 0), dtype=float) + # create context - create_context(m, d, grad=True) + create_context(m, d, h, hfactor, grad=True) # search = -Mgrad wp.launch( @@ -2077,16 +2108,10 @@ def _solve(m: types.Model, d: types.Data): # becomes zero and all worlds are marked as converged to avoid an infinite loop. # note: we only launch the iteration kernel if everything is not done d.nsolving.fill_(d.nworld) - wp.capture_while( - d.nsolving, - while_body=_solver_iteration, - m=m, - d=d, - step_size_cost=step_size_cost, - ) + wp.capture_while(d.nsolving, while_body=_solver_iteration, m=m, d=d, h=h, hfactor=hfactor, step_size_cost=step_size_cost) else: # This branch is mostly for when JAX is used as it is currently not compatible # with CUDA graph conditional. # It should be removed when JAX becomes compatible. for _ in range(m.opt.iterations): - _solver_iteration(m, d, step_size_cost) + _solver_iteration(m, d, h, hfactor, step_size_cost) 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 f4ec1e94..43fb8c37 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -116,10 +116,10 @@ def mul_m_dense(tile: TileSet, check_skip: bool): return dofid = adr[nodeid] - qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid)) - vec_tile = wp.tile_load(vec[worldid], shape=(TILE_SIZE, 1), offset=(dofid, 0)) + qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False) + vec_tile = wp.tile_load(vec[worldid], shape=(TILE_SIZE, 1), offset=(dofid, 0), bounds_check=False) res_tile = wp.tile_matmul(qM_tile, vec_tile) - wp.tile_store(res[worldid], res_tile, offset=(dofid, 0)) + wp.tile_store(res[worldid], res_tile, offset=(dofid, 0), bounds_check=False) return _mul_m_dense 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 1da31efd..6ff16d70 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -58,8 +58,10 @@ class BlockDim: cholesky_factorize_solve: int = 32 # solver update_gradient_cholesky: int = 64 + update_gradient_cholesky_blocked: int = 32 update_gradient_JTDAJ_sparse: int = 64 update_gradient_JTDAJ_dense: int = 96 + linesearch_iterative: int = 64 # support mul_m_dense: int = 32 @@ -502,13 +504,15 @@ class EqType(enum.IntEnum): JOINT: couple the values of two scalar joints with cubic WELD: fix relative position and orientation of two bodies TENDON: couple the lengths of two tendons with cubic + FLEX: couple the edge lengths of a flex """ CONNECT = mujoco.mjtEq.mjEQ_CONNECT WELD = mujoco.mjtEq.mjEQ_WELD JOINT = mujoco.mjtEq.mjEQ_JOINT TENDON = mujoco.mjtEq.mjEQ_TENDON - # unsupported: FLEX, DISTANCE + FLEX = mujoco.mjtEq.mjEQ_FLEX + # unsupported: DISTANCE class WrapType(enum.IntEnum): @@ -638,7 +642,6 @@ class Option: Attributes: timestep: simulation timestep - impratio: ratio of friction-to-normal contact impedance tolerance: main solver tolerance ls_tolerance: CG/Newton linesearch tolerance ccd_tolerance: convex collision detection tolerance @@ -659,6 +662,7 @@ class Option: sdf_iterations: max number of iterations for gradient descent warp only fields: + impratio_invsqrt: ratio of friction-to-normal contact impedance (stored as inverse square root) is_sparse: whether to use sparse representations ls_parallel: evaluate engine solver step sizes in parallel ls_parallel_min_step: minimum step size for solver linesearch @@ -674,7 +678,6 @@ class Option: """ timestep: array("*", float) - impratio: array("*", float) tolerance: array("*", float) ls_tolerance: array("*", float) ccd_tolerance: array("*", float) @@ -694,6 +697,7 @@ class Option: sdf_initpoints: int sdf_iterations: int # warp only fields: + impratio_invsqrt: array("*", float) is_sparse: bool ls_parallel: bool ls_parallel_min_step: float @@ -883,6 +887,9 @@ class Model: flex_vertadr: first vertex address (nflex,) flex_vertnum: number of vertices (nflex,) flex_edgeadr: first edge address (nflex,) + flex_edgenum: number of edges (nflex,) + flex_elemadr: first element address (nflex,) + flex_elemnum: number of elements (nflex,) flex_elemedgeadr: first element address (nflex,) flex_vertbodyid: vertex body ids (nflexvert,) flex_edge: edge vertex ids (2 per edge) (nflexedge, 2) @@ -890,6 +897,7 @@ class Model: flex_elem: element vertex ids (dim+1 per elem) (nflexelemdata,) flex_elemedge: element edge ids (nflexelemedge,) flexedge_length0: edge lengths in qpos0 (nflexedge,) + flexedge_invweight0: inv. inertia for the edge (nflexedge,) flex_stiffness: finite element stiffness matrix (nflexelem, 21) flex_bending: bending stiffness (nflexedge, 17) flex_damping: Rayleigh's damping coefficient (nflex,) @@ -998,6 +1006,7 @@ class Model: mapM2M: index mapping from M (legacy) to M (CSR) (nC) warp only fields: + nv_pad: number of degrees of freedom + padding nacttrnbody: number of actuators with body transmission nsensorcollision: number of unique collisions for geom distance sensors @@ -1031,6 +1040,7 @@ class Model: eq_wld_adr: eq_* addresses of type `WELD` eq_jnt_adr: eq_* addresses of type `JOINT` eq_ten_adr: eq_* addresses of type `TENDON` + eq_flex_adr: eq * addresses of type `FLEX tendon_jnt_adr: joint tendon address tendon_site_pair_adr: site pair tendon address tendon_geom_adr: geom tendon address @@ -1225,6 +1235,9 @@ class Model: flex_vertadr: array("nflex", int) flex_vertnum: array("nflex", int) flex_edgeadr: array("nflex", int) + flex_edgenum: array("nflex", int) + flex_elemadr: array("nflex", int) + flex_elemnum: array("nflex", int) flex_elemedgeadr: array("nflex", int) flex_vertbodyid: array("nflexvert", int) flex_edge: array("nflexedge", wp.vec2i) @@ -1232,6 +1245,7 @@ class Model: flex_elem: array("nflexelemdata", int) flex_elemedge: array("nflexelemedge", int) flexedge_length0: array("nflexedge", float) + flexedge_invweight0: array("nflexedge", float) flex_stiffness: array("nflexelem", 21, float) flex_bending: array("nflexedge", 17, float) flex_damping: array("nflex", float) @@ -1339,6 +1353,7 @@ class Model: M_colind: array("nC", int) mapM2M: array("nC", int) # warp only fields: + nv_pad: int nacttrnbody: int nsensorcollision: int nsensortaxel: int @@ -1367,6 +1382,7 @@ class Model: eq_wld_adr: wp.array(dtype=int) eq_jnt_adr: wp.array(dtype=int) eq_ten_adr: wp.array(dtype=int) + eq_flex_adr: wp.array(dtype=int) tendon_jnt_adr: wp.array(dtype=int) tendon_site_pair_adr: wp.array(dtype=int) tendon_geom_adr: wp.array(dtype=int) @@ -1476,11 +1492,9 @@ class Constraint: 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) - cholesky_L_tmp: temporary for Cholesky factor (nworld, nv, nv) - cholesky_y_tmp: temporary for Cholesky solve (nworld, nv + grad: gradient of master cost (nworld, nv_pad) grad_dot: dot(grad, grad) (nworld,) - Mgrad: M / grad (nworld, nv) + Mgrad: M / grad (nworld, nv_pad) search: linesearch vector (nworld, nv) search_dot: dot(search, search) (nworld,) gauss: Gauss Cost (nworld,) @@ -1491,7 +1505,6 @@ class Constraint: 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_pad, nv_pad) alpha: line search step size (nworld,) prev_grad: previous grad (nworld, nv) prev_Mgrad: previous Mgrad (nworld, nv) @@ -1511,11 +1524,9 @@ class Constraint: force: array("nworld", "njmax", float) Jaref: array("nworld", "njmax", float) Ma: array("nworld", "nv", float) - grad: array("nworld", "nv", float) - cholesky_L_tmp: array("nworld", "nv", "nv", float) - cholesky_y_tmp: array("nworld", "nv", float) + grad: array("nworld", "nv_pad", float) grad_dot: array("nworld", float) - Mgrad: array("nworld", "nv", float) + Mgrad: array("nworld", "nv_pad", float) search: array("nworld", "nv", float) search_dot: array("nworld", float) gauss: array("nworld", float) @@ -1526,7 +1537,6 @@ class Constraint: jv: array("nworld", "njmax", float) quad: array("nworld", "njmax", wp.vec3) quad_gauss: array("nworld", wp.vec3) - h: array("nworld", "nv_pad", "nv_pad", float) alpha: array("nworld", float) prev_grad: array("nworld", "nv", float) prev_Mgrad: array("nworld", "nv", float) @@ -1578,6 +1588,7 @@ class Data: cdof: com-based motion axis of each dof (rot:lin) (nworld, nv, 6) cinert: com-based body inertia and mass (nworld, nbody, 10) flexvert_xpos: cartesian flex vertex positions (nworld, nflexvert, 3) + flexedge_J: edge length Jacobian (nworld, nflexedge, nv) flexedge_length: flex edge lengths (nworld, nflexedge, 1) ten_wrapadr: start address of tendon's path (nworld, ntendon) ten_wrapnum: number of wrap points in path (nworld, ntendon) @@ -1629,6 +1640,7 @@ class Data: ne_weld: number of equality weld constraints (nworld,) ne_jnt: number of equality joint constraints (nworld,) ne_ten: number of equality tendon constraints (nworld,) + ne_flex: number of flex edge equality constraints (nworld,) nsolving: number of unconverged worlds (1,) subtree_bodyvel: subtree body velocity (ang, vel) (nworld, nbody, 6) collision_pair: collision pairs from broadphase (naconmax, 2) @@ -1676,6 +1688,7 @@ class Data: cdof: array("nworld", "nv", wp.spatial_vector) cinert: array("nworld", "nbody", vec10) flexvert_xpos: array("nworld", "nflexvert", wp.vec3) + flexedge_J: array("nworld", "nflexedge", "nv", float) flexedge_length: array("nworld", "nflexedge", float) ten_wrapadr: array("nworld", "ntendon", int) ten_wrapnum: array("nworld", "ntendon", int) @@ -1723,6 +1736,7 @@ class Data: ne_weld: array("nworld", int) ne_jnt: array("nworld", int) ne_ten: array("nworld", int) + ne_flex: array("nworld", int) nsolving: array(1, int) subtree_bodyvel: array("nworld", "nbody", wp.spatial_vector) 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 5a071be0..1544ec63 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,8 @@ import inspect from typing import Callable, Optional import warp as wp -from warp.context import Module -from warp.context import get_module +from warp._src.context import Module +from warp._src.context import get_module _STACK = None @@ -220,8 +220,8 @@ def cache_kernel(func): def check_toolkit_driver(): - if wp.context.runtime is None: - wp.context.init() + if wp._src.context.runtime is None: + wp._src.context.init() if wp.get_device().is_cuda: - if wp.context.runtime.toolkit_version < (12, 4) or wp.context.runtime.driver_version < (12, 4): + if wp._src.context.runtime.toolkit_version < (12, 4) or wp._src.context.runtime.driver_version < (12, 4): 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 f982afff..7006d4c8 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml @@ -18,13 +18,12 @@ classifiers = [ "Intended Audience :: Science/Research", "Natural Language :: English", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering", ] -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = [ "absl-py", "etils[epath]", @@ -35,7 +34,7 @@ dependencies = [ [[tool.uv.index]] name = "nvidia" -url = "https://pypi.nvidia.com" +url = "https://pypi.nvidia.com/" explicit = true [[tool.uv.index]] @@ -56,6 +55,8 @@ dev = [ "ruff", "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", ] # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp cpu = [ diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index c5787fe1..e1702192 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.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 _forward_shim( # Model @@ -104,7 +105,6 @@ def _forward_shim( cam_resolution: wp.array(dtype=wp.vec2i), cam_sensorsize: wp.array(dtype=wp.vec2), cam_targetbodyid: wp.array(dtype=int), - collision_sensor_adr: wp.array(dtype=int), dof_Madr: wp.array(dtype=int), dof_armature: wp.array2d(dtype=float), dof_bodyid: wp.array(dtype=int), @@ -119,6 +119,7 @@ def _forward_shim( dof_tri_row: wp.array(dtype=int), eq_connect_adr: wp.array(dtype=int), eq_data: wp.array2d(dtype=mjwp_types.vec11), + eq_flex_adr: wp.array(dtype=int), eq_jnt_adr: wp.array(dtype=int), eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), @@ -133,12 +134,16 @@ def _forward_shim( flex_edge: wp.array(dtype=wp.vec2i), flex_edgeadr: wp.array(dtype=int), flex_edgeflap: wp.array(dtype=wp.vec2i), + flex_edgenum: wp.array(dtype=int), flex_elem: wp.array(dtype=int), + flex_elemadr: wp.array(dtype=int), flex_elemedge: wp.array(dtype=int), flex_elemedgeadr: wp.array(dtype=int), + flex_elemnum: wp.array(dtype=int), flex_stiffness: wp.array2d(dtype=float), flex_vertadr: wp.array(dtype=int), flex_vertbodyid: wp.array(dtype=int), + flexedge_invweight0: wp.array(dtype=float), flexedge_length0: wp.array(dtype=float), geom_aabb: wp.array3d(dtype=wp.vec3), geom_bodyid: wp.array(dtype=int), @@ -213,13 +218,13 @@ def _forward_shim( mesh_vert: wp.array(dtype=wp.vec3), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), - mocap_bodyid: wp.array(dtype=int), nC: int, na: int, nacttrnbody: int, nbody: int, ncam: int, neq: int, + nflex: int, nflexedge: int, nflexelem: int, nflexvert: int, @@ -232,7 +237,6 @@ def _forward_shim( nmaxpolygon: int, nmaxpyramid: int, nmeshface: int, - nmocap: int, nrangefinder: int, nsensorcollision: int, nsensorcontact: int, @@ -241,6 +245,7 @@ def _forward_shim( ntendon: int, nu: int, nv: int, + nv_pad: int, nwrap: int, nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array(dtype=wp.vec2i), @@ -270,7 +275,6 @@ def _forward_shim( sensor_acc_adr: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), sensor_adr_to_contact_adr: wp.array(dtype=int), - sensor_collision_start_adr: wp.array(dtype=int), sensor_contact_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), sensor_datatype: wp.array(dtype=int), @@ -341,7 +345,7 @@ def _forward_shim( opt__graph_conditional: bool, opt__gravity: wp.array(dtype=wp.vec3), opt__has_fluid: bool, - opt__impratio: wp.array(dtype=float), + opt__impratio_invsqrt: wp.array(dtype=float), opt__is_sparse: bool, opt__iterations: int, opt__ls_iterations: int, @@ -383,6 +387,7 @@ def _forward_shim( cvel: wp.array2d(dtype=wp.spatial_vector), energy: wp.array(dtype=wp.vec2), eq_active: wp.array2d(dtype=bool), + flexedge_J: wp.array3d(dtype=float), flexedge_length: wp.array2d(dtype=float), flexedge_velocity: wp.array2d(dtype=float), flexvert_xpos: wp.array2d(dtype=wp.vec3), @@ -396,6 +401,7 @@ def _forward_shim( ncollision: wp.array(dtype=int), ne: wp.array(dtype=int), ne_connect: wp.array(dtype=int), + ne_flex: wp.array(dtype=int), ne_jnt: wp.array(dtype=int), ne_ten: wp.array(dtype=int), ne_weld: wp.array(dtype=int), @@ -467,8 +473,6 @@ def _forward_shim( efc__alpha: wp.array(dtype=float), efc__aref: wp.array2d(dtype=float), efc__beta: wp.array(dtype=float), - efc__cholesky_L_tmp: wp.array3d(dtype=float), - efc__cholesky_y_tmp: wp.array2d(dtype=float), efc__cost: wp.array(dtype=float), efc__done: wp.array(dtype=bool), efc__force: wp.array2d(dtype=float), @@ -476,7 +480,6 @@ def _forward_shim( efc__gauss: wp.array(dtype=float), efc__grad: wp.array2d(dtype=float), efc__grad_dot: wp.array(dtype=float), - efc__h: wp.array3d(dtype=float), efc__id: wp.array2d(dtype=int), efc__jv: wp.array2d(dtype=float), efc__margin: wp.array2d(dtype=float), @@ -555,7 +558,6 @@ def _forward_shim( _m.cam_resolution = cam_resolution _m.cam_sensorsize = cam_sensorsize _m.cam_targetbodyid = cam_targetbodyid - _m.collision_sensor_adr = collision_sensor_adr _m.dof_Madr = dof_Madr _m.dof_armature = dof_armature _m.dof_bodyid = dof_bodyid @@ -570,6 +572,7 @@ def _forward_shim( _m.dof_tri_row = dof_tri_row _m.eq_connect_adr = eq_connect_adr _m.eq_data = eq_data + _m.eq_flex_adr = eq_flex_adr _m.eq_jnt_adr = eq_jnt_adr _m.eq_obj1id = eq_obj1id _m.eq_obj2id = eq_obj2id @@ -584,12 +587,16 @@ def _forward_shim( _m.flex_edge = flex_edge _m.flex_edgeadr = flex_edgeadr _m.flex_edgeflap = flex_edgeflap + _m.flex_edgenum = flex_edgenum _m.flex_elem = flex_elem + _m.flex_elemadr = flex_elemadr _m.flex_elemedge = flex_elemedge _m.flex_elemedgeadr = flex_elemedgeadr + _m.flex_elemnum = flex_elemnum _m.flex_stiffness = flex_stiffness _m.flex_vertadr = flex_vertadr _m.flex_vertbodyid = flex_vertbodyid + _m.flexedge_invweight0 = flexedge_invweight0 _m.flexedge_length0 = flexedge_length0 _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid @@ -664,13 +671,13 @@ def _forward_shim( _m.mesh_vert = mesh_vert _m.mesh_vertadr = mesh_vertadr _m.mesh_vertnum = mesh_vertnum - _m.mocap_bodyid = mocap_bodyid _m.nC = nC _m.na = na _m.nacttrnbody = nacttrnbody _m.nbody = nbody _m.ncam = ncam _m.neq = neq + _m.nflex = nflex _m.nflexedge = nflexedge _m.nflexelem = nflexelem _m.nflexvert = nflexvert @@ -683,7 +690,6 @@ def _forward_shim( _m.nmaxpolygon = nmaxpolygon _m.nmaxpyramid = nmaxpyramid _m.nmeshface = nmeshface - _m.nmocap = nmocap _m.nrangefinder = nrangefinder _m.nsensorcollision = nsensorcollision _m.nsensorcontact = nsensorcontact @@ -692,6 +698,7 @@ def _forward_shim( _m.ntendon = ntendon _m.nu = nu _m.nv = nv + _m.nv_pad = nv_pad _m.nwrap = nwrap _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered _m.nxn_pairid = nxn_pairid @@ -711,7 +718,7 @@ def _forward_shim( _m.opt.graph_conditional = opt__graph_conditional _m.opt.gravity = opt__gravity _m.opt.has_fluid = opt__has_fluid - _m.opt.impratio = opt__impratio + _m.opt.impratio_invsqrt = opt__impratio_invsqrt _m.opt.is_sparse = opt__is_sparse _m.opt.iterations = opt__iterations _m.opt.ls_iterations = opt__ls_iterations @@ -749,7 +756,6 @@ def _forward_shim( _m.sensor_acc_adr = sensor_acc_adr _m.sensor_adr = sensor_adr _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr - _m.sensor_collision_start_adr = sensor_collision_start_adr _m.sensor_contact_adr = sensor_contact_adr _m.sensor_cutoff = sensor_cutoff _m.sensor_datatype = sensor_datatype @@ -851,8 +857,6 @@ def _forward_shim( _d.efc.alpha = efc__alpha _d.efc.aref = efc__aref _d.efc.beta = efc__beta - _d.efc.cholesky_L_tmp = efc__cholesky_L_tmp - _d.efc.cholesky_y_tmp = efc__cholesky_y_tmp _d.efc.cost = efc__cost _d.efc.done = efc__done _d.efc.force = efc__force @@ -860,7 +864,6 @@ def _forward_shim( _d.efc.gauss = efc__gauss _d.efc.grad = efc__grad _d.efc.grad_dot = efc__grad_dot - _d.efc.h = efc__h _d.efc.id = efc__id _d.efc.jv = efc__jv _d.efc.margin = efc__margin @@ -878,6 +881,7 @@ def _forward_shim( _d.efc.vel = efc__vel _d.energy = energy _d.eq_active = eq_active + _d.flexedge_J = flexedge_J _d.flexedge_length = flexedge_length _d.flexedge_velocity = flexedge_velocity _d.flexvert_xpos = flexvert_xpos @@ -892,6 +896,7 @@ def _forward_shim( _d.ncollision = ncollision _d.ne = ne _d.ne_connect = ne_connect + _d.ne_flex = ne_flex _d.ne_jnt = ne_jnt _d.ne_ten = ne_ten _d.ne_weld = ne_weld @@ -970,6 +975,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'cvel': d.cvel.shape, 'energy': d._impl.energy.shape, 'eq_active': d.eq_active.shape, + 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, 'flexvert_xpos': d._impl.flexvert_xpos.shape, @@ -983,6 +989,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, 'ne_connect': d._impl.ne_connect.shape, + 'ne_flex': d._impl.ne_flex.shape, 'ne_jnt': d._impl.ne_jnt.shape, 'ne_ten': d._impl.ne_ten.shape, 'ne_weld': d._impl.ne_weld.shape, @@ -1054,8 +1061,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__alpha': d._impl.efc__alpha.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__beta': d._impl.efc__beta.shape, - 'efc__cholesky_L_tmp': d._impl.efc__cholesky_L_tmp.shape, - 'efc__cholesky_y_tmp': d._impl.efc__cholesky_y_tmp.shape, 'efc__cost': d._impl.efc__cost.shape, 'efc__done': d._impl.efc__done.shape, 'efc__force': d._impl.efc__force.shape, @@ -1063,7 +1068,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__gauss': d._impl.efc__gauss.shape, 'efc__grad': d._impl.efc__grad.shape, 'efc__grad_dot': d._impl.efc__grad_dot.shape, - 'efc__h': d._impl.efc__h.shape, 'efc__id': d._impl.efc__id.shape, 'efc__jv': d._impl.efc__jv.shape, 'efc__margin': d._impl.efc__margin.shape, @@ -1082,7 +1086,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _forward_shim, - num_outputs=131, + num_outputs=130, output_dims=output_dims, vmap_method=None, in_out_argnames={ @@ -1108,6 +1112,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'cvel', 'energy', 'eq_active', + 'flexedge_J', 'flexedge_length', 'flexedge_velocity', 'flexvert_xpos', @@ -1121,6 +1126,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'ncollision', 'ne', 'ne_connect', + 'ne_flex', 'ne_jnt', 'ne_ten', 'ne_weld', @@ -1192,8 +1198,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__alpha', 'efc__aref', 'efc__beta', - 'efc__cholesky_L_tmp', - 'efc__cholesky_y_tmp', 'efc__cost', 'efc__done', 'efc__force', @@ -1201,7 +1205,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__gauss', 'efc__grad', 'efc__grad_dot', - 'efc__h', 'efc__id', 'efc__jv', 'efc__margin', @@ -1279,7 +1282,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.cam_resolution, m.cam_sensorsize, m.cam_targetbodyid, - m._impl.collision_sensor_adr, m.dof_Madr, m.dof_armature, m.dof_bodyid, @@ -1294,6 +1296,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.dof_tri_row, m._impl.eq_connect_adr, m.eq_data, + m._impl.eq_flex_adr, m._impl.eq_jnt_adr, m.eq_obj1id, m.eq_obj2id, @@ -1308,12 +1311,16 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.flex_edge, m._impl.flex_edgeadr, m._impl.flex_edgeflap, + m._impl.flex_edgenum, m._impl.flex_elem, + m._impl.flex_elemadr, m._impl.flex_elemedge, m._impl.flex_elemedgeadr, + m._impl.flex_elemnum, m._impl.flex_stiffness, m._impl.flex_vertadr, m._impl.flex_vertbodyid, + m._impl.flexedge_invweight0, m._impl.flexedge_length0, m.geom_aabb, m.geom_bodyid, @@ -1388,13 +1395,13 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.mesh_vert, m.mesh_vertadr, m.mesh_vertnum, - m._impl.mocap_bodyid, m.nC, m.na, m._impl.nacttrnbody, m.nbody, m.ncam, m.neq, + m._impl.nflex, m._impl.nflexedge, m._impl.nflexelem, m._impl.nflexvert, @@ -1407,7 +1414,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.nmaxpolygon, m._impl.nmaxpyramid, m.nmeshface, - m.nmocap, m._impl.nrangefinder, m._impl.nsensorcollision, m._impl.nsensorcontact, @@ -1416,6 +1422,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.ntendon, m.nu, m.nv, + m._impl.nv_pad, m.nwrap, m._impl.nxn_geom_pair_filtered, m._impl.nxn_pairid, @@ -1445,7 +1452,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.sensor_acc_adr, m.sensor_adr, m._impl.sensor_adr_to_contact_adr, - m._impl.sensor_collision_start_adr, m._impl.sensor_contact_adr, m.sensor_cutoff, m.sensor_datatype, @@ -1516,7 +1522,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.opt._impl.graph_conditional, m.opt.gravity, m.opt._impl.has_fluid, - m.opt.impratio, + m.opt._impl.impratio_invsqrt, m.opt._impl.is_sparse, m.opt.iterations, m.opt.ls_iterations, @@ -1557,6 +1563,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d.cvel, d._impl.energy, d.eq_active, + d._impl.flexedge_J, d._impl.flexedge_length, d._impl.flexedge_velocity, d._impl.flexvert_xpos, @@ -1570,6 +1577,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.ncollision, d._impl.ne, d._impl.ne_connect, + d._impl.ne_flex, d._impl.ne_jnt, d._impl.ne_ten, d._impl.ne_weld, @@ -1641,8 +1649,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.efc__alpha, d._impl.efc__aref, d._impl.efc__beta, - d._impl.efc__cholesky_L_tmp, - d._impl.efc__cholesky_y_tmp, d._impl.efc__cost, d._impl.efc__done, d._impl.efc__force, @@ -1650,7 +1656,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.efc__gauss, d._impl.efc__grad, d._impl.efc__grad_dot, - d._impl.efc__h, d._impl.efc__id, d._impl.efc__jv, d._impl.efc__margin, @@ -1690,92 +1695,92 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'cvel': out[19], '_impl.energy': out[20], 'eq_active': out[21], - '_impl.flexedge_length': out[22], - '_impl.flexedge_velocity': out[23], - '_impl.flexvert_xpos': out[24], - 'geom_xmat': out[25], - 'geom_xpos': out[26], - '_impl.light_xdir': out[27], - '_impl.light_xpos': out[28], - 'mocap_pos': out[29], - 'mocap_quat': out[30], - '_impl.nacon': out[31], - '_impl.ncollision': out[32], - '_impl.ne': out[33], - '_impl.ne_connect': out[34], - '_impl.ne_jnt': out[35], - '_impl.ne_ten': out[36], - '_impl.ne_weld': out[37], - '_impl.nefc': out[38], - '_impl.nf': out[39], - '_impl.nl': out[40], - '_impl.nsolving': out[41], - '_impl.qLD': out[42], - '_impl.qLDiagInv': out[43], - '_impl.qM': out[44], - 'qacc': out[45], - 'qacc_smooth': out[46], - 'qacc_warmstart': out[47], - 'qfrc_actuator': out[48], - 'qfrc_applied': out[49], - 'qfrc_bias': out[50], - 'qfrc_constraint': out[51], - '_impl.qfrc_damper': out[52], - 'qfrc_fluid': out[53], - 'qfrc_gravcomp': out[54], - 'qfrc_passive': out[55], - 'qfrc_smooth': out[56], - '_impl.qfrc_spring': out[57], - 'qpos': out[58], - 'qvel': out[59], - 'sensordata': out[60], - 'site_xmat': out[61], - 'site_xpos': out[62], - '_impl.solver_niter': out[63], - '_impl.subtree_angmom': out[64], - '_impl.subtree_bodyvel': out[65], - 'subtree_com': out[66], - '_impl.subtree_linvel': out[67], - '_impl.ten_J': out[68], - 'ten_length': out[69], - '_impl.ten_velocity': out[70], - '_impl.ten_wrapadr': out[71], - '_impl.ten_wrapnum': out[72], - 'time': out[73], - '_impl.wrap_obj': out[74], - '_impl.wrap_xpos': out[75], - 'xanchor': out[76], - 'xaxis': out[77], - 'xfrc_applied': out[78], - 'ximat': out[79], - 'xipos': out[80], - 'xmat': out[81], - 'xpos': out[82], - 'xquat': out[83], - '_impl.contact__dim': out[84], - '_impl.contact__dist': out[85], - '_impl.contact__efc_address': out[86], - '_impl.contact__frame': out[87], - '_impl.contact__friction': out[88], - '_impl.contact__geom': out[89], - '_impl.contact__geomcollisionid': out[90], - '_impl.contact__includemargin': out[91], - '_impl.contact__pos': out[92], - '_impl.contact__solimp': out[93], - '_impl.contact__solref': out[94], - '_impl.contact__solreffriction': out[95], - '_impl.contact__type': out[96], - '_impl.contact__worldid': out[97], - '_impl.efc__D': out[98], - '_impl.efc__J': out[99], - '_impl.efc__Jaref': out[100], - '_impl.efc__Ma': out[101], - '_impl.efc__Mgrad': out[102], - '_impl.efc__alpha': out[103], - '_impl.efc__aref': out[104], - '_impl.efc__beta': out[105], - '_impl.efc__cholesky_L_tmp': out[106], - '_impl.efc__cholesky_y_tmp': out[107], + '_impl.flexedge_J': out[22], + '_impl.flexedge_length': out[23], + '_impl.flexedge_velocity': out[24], + '_impl.flexvert_xpos': out[25], + 'geom_xmat': out[26], + 'geom_xpos': out[27], + '_impl.light_xdir': out[28], + '_impl.light_xpos': out[29], + 'mocap_pos': out[30], + 'mocap_quat': out[31], + '_impl.nacon': out[32], + '_impl.ncollision': out[33], + '_impl.ne': out[34], + '_impl.ne_connect': out[35], + '_impl.ne_flex': out[36], + '_impl.ne_jnt': out[37], + '_impl.ne_ten': out[38], + '_impl.ne_weld': out[39], + '_impl.nefc': out[40], + '_impl.nf': out[41], + '_impl.nl': out[42], + '_impl.nsolving': out[43], + '_impl.qLD': out[44], + '_impl.qLDiagInv': out[45], + '_impl.qM': out[46], + 'qacc': out[47], + 'qacc_smooth': out[48], + 'qacc_warmstart': out[49], + 'qfrc_actuator': out[50], + 'qfrc_applied': out[51], + 'qfrc_bias': out[52], + 'qfrc_constraint': out[53], + '_impl.qfrc_damper': out[54], + 'qfrc_fluid': out[55], + 'qfrc_gravcomp': out[56], + 'qfrc_passive': out[57], + 'qfrc_smooth': out[58], + '_impl.qfrc_spring': out[59], + 'qpos': out[60], + 'qvel': out[61], + 'sensordata': out[62], + 'site_xmat': out[63], + 'site_xpos': out[64], + '_impl.solver_niter': out[65], + '_impl.subtree_angmom': out[66], + '_impl.subtree_bodyvel': out[67], + 'subtree_com': out[68], + '_impl.subtree_linvel': out[69], + '_impl.ten_J': out[70], + 'ten_length': out[71], + '_impl.ten_velocity': out[72], + '_impl.ten_wrapadr': out[73], + '_impl.ten_wrapnum': out[74], + 'time': out[75], + '_impl.wrap_obj': out[76], + '_impl.wrap_xpos': out[77], + 'xanchor': out[78], + 'xaxis': out[79], + 'xfrc_applied': out[80], + 'ximat': out[81], + 'xipos': out[82], + 'xmat': out[83], + 'xpos': out[84], + 'xquat': out[85], + '_impl.contact__dim': out[86], + '_impl.contact__dist': out[87], + '_impl.contact__efc_address': out[88], + '_impl.contact__frame': out[89], + '_impl.contact__friction': out[90], + '_impl.contact__geom': out[91], + '_impl.contact__geomcollisionid': out[92], + '_impl.contact__includemargin': out[93], + '_impl.contact__pos': out[94], + '_impl.contact__solimp': out[95], + '_impl.contact__solref': out[96], + '_impl.contact__solreffriction': out[97], + '_impl.contact__type': out[98], + '_impl.contact__worldid': out[99], + '_impl.efc__D': out[100], + '_impl.efc__J': out[101], + '_impl.efc__Jaref': out[102], + '_impl.efc__Ma': out[103], + '_impl.efc__Mgrad': out[104], + '_impl.efc__alpha': out[105], + '_impl.efc__aref': out[106], + '_impl.efc__beta': out[107], '_impl.efc__cost': out[108], '_impl.efc__done': out[109], '_impl.efc__force': out[110], @@ -1783,22 +1788,21 @@ def _forward_jax_impl(m: types.Model, d: types.Data): '_impl.efc__gauss': out[112], '_impl.efc__grad': out[113], '_impl.efc__grad_dot': out[114], - '_impl.efc__h': out[115], - '_impl.efc__id': out[116], - '_impl.efc__jv': out[117], - '_impl.efc__margin': out[118], - '_impl.efc__mv': out[119], - '_impl.efc__pos': out[120], - '_impl.efc__prev_Mgrad': out[121], - '_impl.efc__prev_cost': out[122], - '_impl.efc__prev_grad': out[123], - '_impl.efc__quad': out[124], - '_impl.efc__quad_gauss': out[125], - '_impl.efc__search': out[126], - '_impl.efc__search_dot': out[127], - '_impl.efc__state': out[128], - '_impl.efc__type': out[129], - '_impl.efc__vel': out[130], + '_impl.efc__id': out[115], + '_impl.efc__jv': out[116], + '_impl.efc__margin': out[117], + '_impl.efc__mv': out[118], + '_impl.efc__pos': out[119], + '_impl.efc__prev_Mgrad': out[120], + '_impl.efc__prev_cost': out[121], + '_impl.efc__prev_grad': out[122], + '_impl.efc__quad': out[123], + '_impl.efc__quad_gauss': out[124], + '_impl.efc__search': out[125], + '_impl.efc__search_dot': out[126], + '_impl.efc__state': out[127], + '_impl.efc__type': out[128], + '_impl.efc__vel': out[129], }) return d @@ -1833,6 +1837,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _step_shim( # Model @@ -1895,7 +1900,6 @@ def _step_shim( cam_resolution: wp.array(dtype=wp.vec2i), cam_sensorsize: wp.array(dtype=wp.vec2), cam_targetbodyid: wp.array(dtype=int), - collision_sensor_adr: wp.array(dtype=int), dof_Madr: wp.array(dtype=int), dof_armature: wp.array2d(dtype=float), dof_bodyid: wp.array(dtype=int), @@ -1910,6 +1914,7 @@ def _step_shim( dof_tri_row: wp.array(dtype=int), eq_connect_adr: wp.array(dtype=int), eq_data: wp.array2d(dtype=mjwp_types.vec11), + eq_flex_adr: wp.array(dtype=int), eq_jnt_adr: wp.array(dtype=int), eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), @@ -1924,12 +1929,16 @@ def _step_shim( flex_edge: wp.array(dtype=wp.vec2i), flex_edgeadr: wp.array(dtype=int), flex_edgeflap: wp.array(dtype=wp.vec2i), + flex_edgenum: wp.array(dtype=int), flex_elem: wp.array(dtype=int), + flex_elemadr: wp.array(dtype=int), flex_elemedge: wp.array(dtype=int), flex_elemedgeadr: wp.array(dtype=int), + flex_elemnum: wp.array(dtype=int), flex_stiffness: wp.array2d(dtype=float), flex_vertadr: wp.array(dtype=int), flex_vertbodyid: wp.array(dtype=int), + flexedge_invweight0: wp.array(dtype=float), flexedge_length0: wp.array(dtype=float), geom_aabb: wp.array3d(dtype=wp.vec3), geom_bodyid: wp.array(dtype=int), @@ -2004,7 +2013,6 @@ def _step_shim( mesh_vert: wp.array(dtype=wp.vec3), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), - mocap_bodyid: wp.array(dtype=int), nC: int, nM: int, na: int, @@ -2012,6 +2020,7 @@ def _step_shim( nbody: int, ncam: int, neq: int, + nflex: int, nflexedge: int, nflexelem: int, nflexvert: int, @@ -2024,7 +2033,6 @@ def _step_shim( nmaxpolygon: int, nmaxpyramid: int, nmeshface: int, - nmocap: int, nrangefinder: int, nsensorcollision: int, nsensorcontact: int, @@ -2033,6 +2041,7 @@ def _step_shim( ntendon: int, nu: int, nv: int, + nv_pad: int, nwrap: int, nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array(dtype=wp.vec2i), @@ -2062,7 +2071,6 @@ def _step_shim( sensor_acc_adr: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), sensor_adr_to_contact_adr: wp.array(dtype=int), - sensor_collision_start_adr: wp.array(dtype=int), sensor_contact_adr: wp.array(dtype=int), sensor_cutoff: wp.array(dtype=float), sensor_datatype: wp.array(dtype=int), @@ -2133,7 +2141,7 @@ def _step_shim( opt__graph_conditional: bool, opt__gravity: wp.array(dtype=wp.vec3), opt__has_fluid: bool, - opt__impratio: wp.array(dtype=float), + opt__impratio_invsqrt: wp.array(dtype=float), opt__integrator: int, opt__is_sparse: bool, opt__iterations: int, @@ -2176,6 +2184,7 @@ def _step_shim( cvel: wp.array2d(dtype=wp.spatial_vector), energy: wp.array(dtype=wp.vec2), eq_active: wp.array2d(dtype=bool), + flexedge_J: wp.array3d(dtype=float), flexedge_length: wp.array2d(dtype=float), flexedge_velocity: wp.array2d(dtype=float), flexvert_xpos: wp.array2d(dtype=wp.vec3), @@ -2189,6 +2198,7 @@ def _step_shim( ncollision: wp.array(dtype=int), ne: wp.array(dtype=int), ne_connect: wp.array(dtype=int), + ne_flex: wp.array(dtype=int), ne_jnt: wp.array(dtype=int), ne_ten: wp.array(dtype=int), ne_weld: wp.array(dtype=int), @@ -2260,8 +2270,6 @@ def _step_shim( efc__alpha: wp.array(dtype=float), efc__aref: wp.array2d(dtype=float), efc__beta: wp.array(dtype=float), - efc__cholesky_L_tmp: wp.array3d(dtype=float), - efc__cholesky_y_tmp: wp.array2d(dtype=float), efc__cost: wp.array(dtype=float), efc__done: wp.array(dtype=bool), efc__force: wp.array2d(dtype=float), @@ -2269,7 +2277,6 @@ def _step_shim( efc__gauss: wp.array(dtype=float), efc__grad: wp.array2d(dtype=float), efc__grad_dot: wp.array(dtype=float), - efc__h: wp.array3d(dtype=float), efc__id: wp.array2d(dtype=int), efc__jv: wp.array2d(dtype=float), efc__margin: wp.array2d(dtype=float), @@ -2348,7 +2355,6 @@ def _step_shim( _m.cam_resolution = cam_resolution _m.cam_sensorsize = cam_sensorsize _m.cam_targetbodyid = cam_targetbodyid - _m.collision_sensor_adr = collision_sensor_adr _m.dof_Madr = dof_Madr _m.dof_armature = dof_armature _m.dof_bodyid = dof_bodyid @@ -2363,6 +2369,7 @@ def _step_shim( _m.dof_tri_row = dof_tri_row _m.eq_connect_adr = eq_connect_adr _m.eq_data = eq_data + _m.eq_flex_adr = eq_flex_adr _m.eq_jnt_adr = eq_jnt_adr _m.eq_obj1id = eq_obj1id _m.eq_obj2id = eq_obj2id @@ -2377,12 +2384,16 @@ def _step_shim( _m.flex_edge = flex_edge _m.flex_edgeadr = flex_edgeadr _m.flex_edgeflap = flex_edgeflap + _m.flex_edgenum = flex_edgenum _m.flex_elem = flex_elem + _m.flex_elemadr = flex_elemadr _m.flex_elemedge = flex_elemedge _m.flex_elemedgeadr = flex_elemedgeadr + _m.flex_elemnum = flex_elemnum _m.flex_stiffness = flex_stiffness _m.flex_vertadr = flex_vertadr _m.flex_vertbodyid = flex_vertbodyid + _m.flexedge_invweight0 = flexedge_invweight0 _m.flexedge_length0 = flexedge_length0 _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid @@ -2457,7 +2468,6 @@ def _step_shim( _m.mesh_vert = mesh_vert _m.mesh_vertadr = mesh_vertadr _m.mesh_vertnum = mesh_vertnum - _m.mocap_bodyid = mocap_bodyid _m.nC = nC _m.nM = nM _m.na = na @@ -2465,6 +2475,7 @@ def _step_shim( _m.nbody = nbody _m.ncam = ncam _m.neq = neq + _m.nflex = nflex _m.nflexedge = nflexedge _m.nflexelem = nflexelem _m.nflexvert = nflexvert @@ -2477,7 +2488,6 @@ def _step_shim( _m.nmaxpolygon = nmaxpolygon _m.nmaxpyramid = nmaxpyramid _m.nmeshface = nmeshface - _m.nmocap = nmocap _m.nrangefinder = nrangefinder _m.nsensorcollision = nsensorcollision _m.nsensorcontact = nsensorcontact @@ -2486,6 +2496,7 @@ def _step_shim( _m.ntendon = ntendon _m.nu = nu _m.nv = nv + _m.nv_pad = nv_pad _m.nwrap = nwrap _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered _m.nxn_pairid = nxn_pairid @@ -2505,7 +2516,7 @@ def _step_shim( _m.opt.graph_conditional = opt__graph_conditional _m.opt.gravity = opt__gravity _m.opt.has_fluid = opt__has_fluid - _m.opt.impratio = opt__impratio + _m.opt.impratio_invsqrt = opt__impratio_invsqrt _m.opt.integrator = opt__integrator _m.opt.is_sparse = opt__is_sparse _m.opt.iterations = opt__iterations @@ -2544,7 +2555,6 @@ def _step_shim( _m.sensor_acc_adr = sensor_acc_adr _m.sensor_adr = sensor_adr _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr - _m.sensor_collision_start_adr = sensor_collision_start_adr _m.sensor_contact_adr = sensor_contact_adr _m.sensor_cutoff = sensor_cutoff _m.sensor_datatype = sensor_datatype @@ -2646,8 +2656,6 @@ def _step_shim( _d.efc.alpha = efc__alpha _d.efc.aref = efc__aref _d.efc.beta = efc__beta - _d.efc.cholesky_L_tmp = efc__cholesky_L_tmp - _d.efc.cholesky_y_tmp = efc__cholesky_y_tmp _d.efc.cost = efc__cost _d.efc.done = efc__done _d.efc.force = efc__force @@ -2655,7 +2663,6 @@ def _step_shim( _d.efc.gauss = efc__gauss _d.efc.grad = efc__grad _d.efc.grad_dot = efc__grad_dot - _d.efc.h = efc__h _d.efc.id = efc__id _d.efc.jv = efc__jv _d.efc.margin = efc__margin @@ -2673,6 +2680,7 @@ def _step_shim( _d.efc.vel = efc__vel _d.energy = energy _d.eq_active = eq_active + _d.flexedge_J = flexedge_J _d.flexedge_length = flexedge_length _d.flexedge_velocity = flexedge_velocity _d.flexvert_xpos = flexvert_xpos @@ -2687,6 +2695,7 @@ def _step_shim( _d.ncollision = ncollision _d.ne = ne _d.ne_connect = ne_connect + _d.ne_flex = ne_flex _d.ne_jnt = ne_jnt _d.ne_ten = ne_ten _d.ne_weld = ne_weld @@ -2765,6 +2774,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'cvel': d.cvel.shape, 'energy': d._impl.energy.shape, 'eq_active': d.eq_active.shape, + 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, 'flexvert_xpos': d._impl.flexvert_xpos.shape, @@ -2778,6 +2788,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, 'ne_connect': d._impl.ne_connect.shape, + 'ne_flex': d._impl.ne_flex.shape, 'ne_jnt': d._impl.ne_jnt.shape, 'ne_ten': d._impl.ne_ten.shape, 'ne_weld': d._impl.ne_weld.shape, @@ -2849,8 +2860,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__alpha': d._impl.efc__alpha.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__beta': d._impl.efc__beta.shape, - 'efc__cholesky_L_tmp': d._impl.efc__cholesky_L_tmp.shape, - 'efc__cholesky_y_tmp': d._impl.efc__cholesky_y_tmp.shape, 'efc__cost': d._impl.efc__cost.shape, 'efc__done': d._impl.efc__done.shape, 'efc__force': d._impl.efc__force.shape, @@ -2858,7 +2867,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__gauss': d._impl.efc__gauss.shape, 'efc__grad': d._impl.efc__grad.shape, 'efc__grad_dot': d._impl.efc__grad_dot.shape, - 'efc__h': d._impl.efc__h.shape, 'efc__id': d._impl.efc__id.shape, 'efc__jv': d._impl.efc__jv.shape, 'efc__margin': d._impl.efc__margin.shape, @@ -2877,7 +2885,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _step_shim, - num_outputs=131, + num_outputs=130, output_dims=output_dims, vmap_method=None, in_out_argnames={ @@ -2903,6 +2911,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'cvel', 'energy', 'eq_active', + 'flexedge_J', 'flexedge_length', 'flexedge_velocity', 'flexvert_xpos', @@ -2916,6 +2925,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'ncollision', 'ne', 'ne_connect', + 'ne_flex', 'ne_jnt', 'ne_ten', 'ne_weld', @@ -2987,8 +2997,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__alpha', 'efc__aref', 'efc__beta', - 'efc__cholesky_L_tmp', - 'efc__cholesky_y_tmp', 'efc__cost', 'efc__done', 'efc__force', @@ -2996,7 +3004,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__gauss', 'efc__grad', 'efc__grad_dot', - 'efc__h', 'efc__id', 'efc__jv', 'efc__margin', @@ -3074,7 +3081,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.cam_resolution, m.cam_sensorsize, m.cam_targetbodyid, - m._impl.collision_sensor_adr, m.dof_Madr, m.dof_armature, m.dof_bodyid, @@ -3089,6 +3095,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.dof_tri_row, m._impl.eq_connect_adr, m.eq_data, + m._impl.eq_flex_adr, m._impl.eq_jnt_adr, m.eq_obj1id, m.eq_obj2id, @@ -3103,12 +3110,16 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.flex_edge, m._impl.flex_edgeadr, m._impl.flex_edgeflap, + m._impl.flex_edgenum, m._impl.flex_elem, + m._impl.flex_elemadr, m._impl.flex_elemedge, m._impl.flex_elemedgeadr, + m._impl.flex_elemnum, m._impl.flex_stiffness, m._impl.flex_vertadr, m._impl.flex_vertbodyid, + m._impl.flexedge_invweight0, m._impl.flexedge_length0, m.geom_aabb, m.geom_bodyid, @@ -3183,7 +3194,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.mesh_vert, m.mesh_vertadr, m.mesh_vertnum, - m._impl.mocap_bodyid, m.nC, m.nM, m.na, @@ -3191,6 +3201,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.nbody, m.ncam, m.neq, + m._impl.nflex, m._impl.nflexedge, m._impl.nflexelem, m._impl.nflexvert, @@ -3203,7 +3214,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.nmaxpolygon, m._impl.nmaxpyramid, m.nmeshface, - m.nmocap, m._impl.nrangefinder, m._impl.nsensorcollision, m._impl.nsensorcontact, @@ -3212,6 +3222,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.ntendon, m.nu, m.nv, + m._impl.nv_pad, m.nwrap, m._impl.nxn_geom_pair_filtered, m._impl.nxn_pairid, @@ -3241,7 +3252,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.sensor_acc_adr, m.sensor_adr, m._impl.sensor_adr_to_contact_adr, - m._impl.sensor_collision_start_adr, m._impl.sensor_contact_adr, m.sensor_cutoff, m.sensor_datatype, @@ -3312,7 +3322,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.opt._impl.graph_conditional, m.opt.gravity, m.opt._impl.has_fluid, - m.opt.impratio, + m.opt._impl.impratio_invsqrt, m.opt.integrator, m.opt._impl.is_sparse, m.opt.iterations, @@ -3354,6 +3364,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): d.cvel, d._impl.energy, d.eq_active, + d._impl.flexedge_J, d._impl.flexedge_length, d._impl.flexedge_velocity, d._impl.flexvert_xpos, @@ -3367,6 +3378,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.ncollision, d._impl.ne, d._impl.ne_connect, + d._impl.ne_flex, d._impl.ne_jnt, d._impl.ne_ten, d._impl.ne_weld, @@ -3438,8 +3450,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.efc__alpha, d._impl.efc__aref, d._impl.efc__beta, - d._impl.efc__cholesky_L_tmp, - d._impl.efc__cholesky_y_tmp, d._impl.efc__cost, d._impl.efc__done, d._impl.efc__force, @@ -3447,7 +3457,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.efc__gauss, d._impl.efc__grad, d._impl.efc__grad_dot, - d._impl.efc__h, d._impl.efc__id, d._impl.efc__jv, d._impl.efc__margin, @@ -3487,92 +3496,92 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'cvel': out[19], '_impl.energy': out[20], 'eq_active': out[21], - '_impl.flexedge_length': out[22], - '_impl.flexedge_velocity': out[23], - '_impl.flexvert_xpos': out[24], - 'geom_xmat': out[25], - 'geom_xpos': out[26], - '_impl.light_xdir': out[27], - '_impl.light_xpos': out[28], - 'mocap_pos': out[29], - 'mocap_quat': out[30], - '_impl.nacon': out[31], - '_impl.ncollision': out[32], - '_impl.ne': out[33], - '_impl.ne_connect': out[34], - '_impl.ne_jnt': out[35], - '_impl.ne_ten': out[36], - '_impl.ne_weld': out[37], - '_impl.nefc': out[38], - '_impl.nf': out[39], - '_impl.nl': out[40], - '_impl.nsolving': out[41], - '_impl.qLD': out[42], - '_impl.qLDiagInv': out[43], - '_impl.qM': out[44], - 'qacc': out[45], - 'qacc_smooth': out[46], - 'qacc_warmstart': out[47], - 'qfrc_actuator': out[48], - 'qfrc_applied': out[49], - 'qfrc_bias': out[50], - 'qfrc_constraint': out[51], - '_impl.qfrc_damper': out[52], - 'qfrc_fluid': out[53], - 'qfrc_gravcomp': out[54], - 'qfrc_passive': out[55], - 'qfrc_smooth': out[56], - '_impl.qfrc_spring': out[57], - 'qpos': out[58], - 'qvel': out[59], - 'sensordata': out[60], - 'site_xmat': out[61], - 'site_xpos': out[62], - '_impl.solver_niter': out[63], - '_impl.subtree_angmom': out[64], - '_impl.subtree_bodyvel': out[65], - 'subtree_com': out[66], - '_impl.subtree_linvel': out[67], - '_impl.ten_J': out[68], - 'ten_length': out[69], - '_impl.ten_velocity': out[70], - '_impl.ten_wrapadr': out[71], - '_impl.ten_wrapnum': out[72], - 'time': out[73], - '_impl.wrap_obj': out[74], - '_impl.wrap_xpos': out[75], - 'xanchor': out[76], - 'xaxis': out[77], - 'xfrc_applied': out[78], - 'ximat': out[79], - 'xipos': out[80], - 'xmat': out[81], - 'xpos': out[82], - 'xquat': out[83], - '_impl.contact__dim': out[84], - '_impl.contact__dist': out[85], - '_impl.contact__efc_address': out[86], - '_impl.contact__frame': out[87], - '_impl.contact__friction': out[88], - '_impl.contact__geom': out[89], - '_impl.contact__geomcollisionid': out[90], - '_impl.contact__includemargin': out[91], - '_impl.contact__pos': out[92], - '_impl.contact__solimp': out[93], - '_impl.contact__solref': out[94], - '_impl.contact__solreffriction': out[95], - '_impl.contact__type': out[96], - '_impl.contact__worldid': out[97], - '_impl.efc__D': out[98], - '_impl.efc__J': out[99], - '_impl.efc__Jaref': out[100], - '_impl.efc__Ma': out[101], - '_impl.efc__Mgrad': out[102], - '_impl.efc__alpha': out[103], - '_impl.efc__aref': out[104], - '_impl.efc__beta': out[105], - '_impl.efc__cholesky_L_tmp': out[106], - '_impl.efc__cholesky_y_tmp': out[107], + '_impl.flexedge_J': out[22], + '_impl.flexedge_length': out[23], + '_impl.flexedge_velocity': out[24], + '_impl.flexvert_xpos': out[25], + 'geom_xmat': out[26], + 'geom_xpos': out[27], + '_impl.light_xdir': out[28], + '_impl.light_xpos': out[29], + 'mocap_pos': out[30], + 'mocap_quat': out[31], + '_impl.nacon': out[32], + '_impl.ncollision': out[33], + '_impl.ne': out[34], + '_impl.ne_connect': out[35], + '_impl.ne_flex': out[36], + '_impl.ne_jnt': out[37], + '_impl.ne_ten': out[38], + '_impl.ne_weld': out[39], + '_impl.nefc': out[40], + '_impl.nf': out[41], + '_impl.nl': out[42], + '_impl.nsolving': out[43], + '_impl.qLD': out[44], + '_impl.qLDiagInv': out[45], + '_impl.qM': out[46], + 'qacc': out[47], + 'qacc_smooth': out[48], + 'qacc_warmstart': out[49], + 'qfrc_actuator': out[50], + 'qfrc_applied': out[51], + 'qfrc_bias': out[52], + 'qfrc_constraint': out[53], + '_impl.qfrc_damper': out[54], + 'qfrc_fluid': out[55], + 'qfrc_gravcomp': out[56], + 'qfrc_passive': out[57], + 'qfrc_smooth': out[58], + '_impl.qfrc_spring': out[59], + 'qpos': out[60], + 'qvel': out[61], + 'sensordata': out[62], + 'site_xmat': out[63], + 'site_xpos': out[64], + '_impl.solver_niter': out[65], + '_impl.subtree_angmom': out[66], + '_impl.subtree_bodyvel': out[67], + 'subtree_com': out[68], + '_impl.subtree_linvel': out[69], + '_impl.ten_J': out[70], + 'ten_length': out[71], + '_impl.ten_velocity': out[72], + '_impl.ten_wrapadr': out[73], + '_impl.ten_wrapnum': out[74], + 'time': out[75], + '_impl.wrap_obj': out[76], + '_impl.wrap_xpos': out[77], + 'xanchor': out[78], + 'xaxis': out[79], + 'xfrc_applied': out[80], + 'ximat': out[81], + 'xipos': out[82], + 'xmat': out[83], + 'xpos': out[84], + 'xquat': out[85], + '_impl.contact__dim': out[86], + '_impl.contact__dist': out[87], + '_impl.contact__efc_address': out[88], + '_impl.contact__frame': out[89], + '_impl.contact__friction': out[90], + '_impl.contact__geom': out[91], + '_impl.contact__geomcollisionid': out[92], + '_impl.contact__includemargin': out[93], + '_impl.contact__pos': out[94], + '_impl.contact__solimp': out[95], + '_impl.contact__solref': out[96], + '_impl.contact__solreffriction': out[97], + '_impl.contact__type': out[98], + '_impl.contact__worldid': out[99], + '_impl.efc__D': out[100], + '_impl.efc__J': out[101], + '_impl.efc__Jaref': out[102], + '_impl.efc__Ma': out[103], + '_impl.efc__Mgrad': out[104], + '_impl.efc__alpha': out[105], + '_impl.efc__aref': out[106], + '_impl.efc__beta': out[107], '_impl.efc__cost': out[108], '_impl.efc__done': out[109], '_impl.efc__force': out[110], @@ -3580,22 +3589,21 @@ def _step_jax_impl(m: types.Model, d: types.Data): '_impl.efc__gauss': out[112], '_impl.efc__grad': out[113], '_impl.efc__grad_dot': out[114], - '_impl.efc__h': out[115], - '_impl.efc__id': out[116], - '_impl.efc__jv': out[117], - '_impl.efc__margin': out[118], - '_impl.efc__mv': out[119], - '_impl.efc__pos': out[120], - '_impl.efc__prev_Mgrad': out[121], - '_impl.efc__prev_cost': out[122], - '_impl.efc__prev_grad': out[123], - '_impl.efc__quad': out[124], - '_impl.efc__quad_gauss': out[125], - '_impl.efc__search': out[126], - '_impl.efc__search_dot': out[127], - '_impl.efc__state': out[128], - '_impl.efc__type': out[129], - '_impl.efc__vel': out[130], + '_impl.efc__id': out[115], + '_impl.efc__jv': out[116], + '_impl.efc__margin': out[117], + '_impl.efc__mv': out[118], + '_impl.efc__pos': out[119], + '_impl.efc__prev_Mgrad': out[120], + '_impl.efc__prev_cost': out[121], + '_impl.efc__prev_grad': out[122], + '_impl.efc__quad': out[123], + '_impl.efc__quad_gauss': out[124], + '_impl.efc__search': out[125], + '_impl.efc__search_dot': out[126], + '_impl.efc__state': out[127], + '_impl.efc__type': out[128], + '_impl.efc__vel': out[129], }) return d diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 1412d34e..98c86e14 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -101,6 +101,7 @@ class ForwardTest(parameterized.TestCase): m = test_util.load_test_file(xml) m.opt.iterations = 10 m.opt.ls_iterations = 10 + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE mx = mjx.put_model(m, impl='warp') d = mujoco.MjData(m) @@ -150,11 +151,8 @@ class ForwardTest(parameterized.TestCase): tu.assert_attr_eq(dx._impl, d, 'wrap_obj') tu.assert_attr_eq(dx._impl, d, 'crb') - if not mx.opt._impl.is_sparse: - qm = np.zeros((m.nv, m.nv)) - mujoco.mj_fullM(m, qm, d.qM) - else: - qm = d.qM + qm = np.zeros((m.nv, m.nv)) + mujoco.mj_fullM(m, qm, d.qM) # mjwarp adds padding to qM tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM') # qLD is fused in a cholesky factorize and solve, and not written to. diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index b1aaf5f5..f7a7fb70 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -46,7 +46,6 @@ _e = mjwarp.Constraint( def _kinematics_shim( # Model nworld: int, - body_dofadr: wp.array(dtype=int), body_ipos: wp.array2d(dtype=wp.vec3), body_iquat: wp.array2d(dtype=wp.quat), body_jntadr: wp.array(dtype=int), @@ -58,9 +57,6 @@ def _kinematics_shim( body_rootid: wp.array(dtype=int), body_tree: tuple[wp.array(dtype=int), ...], body_weldid: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_vertadr: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), geom_pos: wp.array2d(dtype=wp.vec3), geom_quat: wp.array2d(dtype=wp.quat), @@ -68,26 +64,18 @@ def _kinematics_shim( jnt_pos: wp.array2d(dtype=wp.vec3), jnt_qposadr: wp.array(dtype=int), jnt_type: wp.array(dtype=int), - mocap_bodyid: wp.array(dtype=int), - nflexedge: int, - nflexvert: int, ngeom: int, - nmocap: int, nsite: int, qpos0: wp.array2d(dtype=float), site_bodyid: wp.array(dtype=int), site_pos: wp.array2d(dtype=wp.vec3), site_quat: wp.array2d(dtype=wp.quat), # Data - flexedge_length: wp.array2d(dtype=float), - flexedge_velocity: wp.array2d(dtype=float), - flexvert_xpos: wp.array2d(dtype=wp.vec3), geom_xmat: wp.array2d(dtype=wp.mat33), geom_xpos: wp.array2d(dtype=wp.vec3), mocap_pos: wp.array2d(dtype=wp.vec3), mocap_quat: wp.array2d(dtype=wp.quat), qpos: wp.array2d(dtype=float), - qvel: wp.array2d(dtype=float), site_xmat: wp.array2d(dtype=wp.mat33), site_xpos: wp.array2d(dtype=wp.vec3), xanchor: wp.array2d(dtype=wp.vec3), @@ -102,7 +90,6 @@ def _kinematics_shim( _m.opt = _o _d.efc = _e _d.contact = _c - _m.body_dofadr = body_dofadr _m.body_ipos = body_ipos _m.body_iquat = body_iquat _m.body_jntadr = body_jntadr @@ -114,9 +101,6 @@ def _kinematics_shim( _m.body_rootid = body_rootid _m.body_tree = body_tree _m.body_weldid = body_weldid - _m.flex_edge = flex_edge - _m.flex_vertadr = flex_vertadr - _m.flex_vertbodyid = flex_vertbodyid _m.geom_bodyid = geom_bodyid _m.geom_pos = geom_pos _m.geom_quat = geom_quat @@ -124,25 +108,17 @@ def _kinematics_shim( _m.jnt_pos = jnt_pos _m.jnt_qposadr = jnt_qposadr _m.jnt_type = jnt_type - _m.mocap_bodyid = mocap_bodyid - _m.nflexedge = nflexedge - _m.nflexvert = nflexvert _m.ngeom = ngeom - _m.nmocap = nmocap _m.nsite = nsite _m.qpos0 = qpos0 _m.site_bodyid = site_bodyid _m.site_pos = site_pos _m.site_quat = site_quat - _d.flexedge_length = flexedge_length - _d.flexedge_velocity = flexedge_velocity - _d.flexvert_xpos = flexvert_xpos _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos _d.mocap_pos = mocap_pos _d.mocap_quat = mocap_quat _d.qpos = qpos - _d.qvel = qvel _d.site_xmat = site_xmat _d.site_xpos = site_xpos _d.xanchor = xanchor @@ -158,15 +134,11 @@ def _kinematics_shim( def _kinematics_jax_impl(m: types.Model, d: types.Data): output_dims = { - 'flexedge_length': d._impl.flexedge_length.shape, - 'flexedge_velocity': d._impl.flexedge_velocity.shape, - 'flexvert_xpos': d._impl.flexvert_xpos.shape, 'geom_xmat': d.geom_xmat.shape, 'geom_xpos': d.geom_xpos.shape, 'mocap_pos': d.mocap_pos.shape, 'mocap_quat': d.mocap_quat.shape, 'qpos': d.qpos.shape, - 'qvel': d.qvel.shape, 'site_xmat': d.site_xmat.shape, 'site_xpos': d.site_xpos.shape, 'xanchor': d.xanchor.shape, @@ -179,19 +151,15 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _kinematics_shim, - num_outputs=18, + num_outputs=14, output_dims=output_dims, vmap_method=None, in_out_argnames={ - 'flexedge_length', - 'flexedge_velocity', - 'flexvert_xpos', 'geom_xmat', 'geom_xpos', 'mocap_pos', 'mocap_quat', 'qpos', - 'qvel', 'site_xmat', 'site_xpos', 'xanchor', @@ -205,7 +173,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): ) out = jf( d.qpos.shape[0], - m.body_dofadr, m.body_ipos, m.body_iquat, m.body_jntadr, @@ -217,9 +184,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): m.body_rootid, m._impl.body_tree, m.body_weldid, - m._impl.flex_edge, - m._impl.flex_vertadr, - m._impl.flex_vertbodyid, m.geom_bodyid, m.geom_pos, m.geom_quat, @@ -227,25 +191,17 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): m.jnt_pos, m.jnt_qposadr, m.jnt_type, - m._impl.mocap_bodyid, - m._impl.nflexedge, - m._impl.nflexvert, m.ngeom, - m.nmocap, m.nsite, m.qpos0, m.site_bodyid, m.site_pos, m.site_quat, - d._impl.flexedge_length, - d._impl.flexedge_velocity, - d._impl.flexvert_xpos, d.geom_xmat, d.geom_xpos, d.mocap_pos, d.mocap_quat, d.qpos, - d.qvel, d.site_xmat, d.site_xpos, d.xanchor, @@ -257,24 +213,20 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): d.xquat, ) d = d.tree_replace({ - '_impl.flexedge_length': out[0], - '_impl.flexedge_velocity': out[1], - '_impl.flexvert_xpos': out[2], - 'geom_xmat': out[3], - 'geom_xpos': out[4], - 'mocap_pos': out[5], - 'mocap_quat': out[6], - 'qpos': out[7], - 'qvel': out[8], - 'site_xmat': out[9], - 'site_xpos': out[10], - 'xanchor': out[11], - 'xaxis': out[12], - 'ximat': out[13], - 'xipos': out[14], - 'xmat': out[15], - 'xpos': out[16], - 'xquat': out[17], + 'geom_xmat': out[0], + 'geom_xpos': out[1], + 'mocap_pos': out[2], + 'mocap_quat': out[3], + 'qpos': out[4], + 'site_xmat': out[5], + 'site_xpos': out[6], + 'xanchor': out[7], + 'xaxis': out[8], + 'ximat': out[9], + 'xipos': out[10], + 'xmat': out[11], + 'xpos': out[12], + 'xquat': out[13], }) return d diff --git a/mjx/mujoco/mjx/warp/smooth_test.py b/mjx/mujoco/mjx/warp/smooth_test.py index c83574fe..b2f867a5 100644 --- a/mjx/mujoco/mjx/warp/smooth_test.py +++ b/mjx/mujoco/mjx/warp/smooth_test.py @@ -73,7 +73,7 @@ class SmoothTest(parameterized.TestCase): _, key1, key2 = jax.random.split(rng, 3) mocap_pos = jax.random.normal(key1, (m.nmocap, 3)) mocap_quat = jax.random.normal(key2, (m.nmocap, 4)) - mocap_quat = math.normalize(mocap_quat) + mocap_quat = math.normalize(mocap_quat, axis=0) dx = dx.replace(qpos=qpos, mocap_pos=mocap_pos, mocap_quat=mocap_quat) dx = jax.jit(smooth.kinematics)(mx, dx) @@ -81,7 +81,7 @@ class SmoothTest(parameterized.TestCase): d.qpos[:] = qpos d.mocap_pos[:] = mocap_pos d.mocap_quat[:] = mocap_quat - mujoco.mj_forward(m, d) + mujoco.mj_kinematics(m, d) tu.assert_attr_eq(d, dx, 'xanchor') tu.assert_attr_eq(d, dx, 'xaxis') @@ -110,9 +110,9 @@ class SmoothTest(parameterized.TestCase): worldids = jp.arange(batch_size) dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids) - fields = ('xanchor', 'xaxis', 'xpos', 'xquat', 'xmat', 'xipos', 'ximat', - 'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat') # fmt: skip - for f in fields: + # don't zero xmat, ximat, xquat, geom_xpos, or geom_xmat + # these fields are precomputed in make_data + for f in ('xanchor', 'xaxis', 'xpos', 'xipos', 'site_xpos', 'site_xmat'): dx_batch = dx_batch.replace(**{f: jp.zeros_like(getattr(dx_batch, f))}) dx_batch = jax.jit(jax.vmap(smooth.kinematics, in_axes=(None, 0)))( @@ -125,7 +125,7 @@ class SmoothTest(parameterized.TestCase): d.qpos[:] = dx.qpos d.mocap_pos[:] = dx.mocap_pos d.mocap_quat[:] = dx.mocap_quat - mujoco.mj_forward(m, d) + mujoco.mj_kinematics(m, d) tu.assert_attr_eq(d, dx, 'xanchor') tu.assert_attr_eq(d, dx, 'xaxis') @@ -154,9 +154,9 @@ class SmoothTest(parameterized.TestCase): worldids = jp.arange(16).reshape((4, 4)) dx_batch = jax.vmap(jax.vmap(functools.partial(tu.make_data, m)))(worldids) - fields = ('xanchor', 'xaxis', 'xpos', 'xquat', 'xmat', 'xipos', 'ximat', - 'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat') # fmt: skip - for f in fields: + # don't zero xmat, ximat, xquat, geom_xpos, or geom_xmat + # these fields are precomputed in make_data + for f in ('xanchor', 'xaxis', 'xpos', 'xipos', 'site_xpos', 'site_xmat'): dx_batch = dx_batch.replace(**{f: jp.zeros_like(getattr(dx_batch, f))}) dx_batch = jax.jit( @@ -208,9 +208,9 @@ class SmoothTest(parameterized.TestCase): worldids = jp.arange(batch_size) dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids) - fields = ('xanchor', 'xaxis', 'xpos', 'xquat', 'xmat', 'xipos', 'ximat', - 'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat') # fmt: skip - for f in fields: + # don't zero xmat, ximat, xquat, geom_xpos, or geom_xmat + # these fields are precomputed in make_data + for f in ('xanchor', 'xaxis', 'xpos', 'xipos', 'site_xpos', 'site_xmat'): dx_batch = dx_batch.replace(**{f: jp.zeros_like(getattr(dx_batch, f))}) dx_batch = jax.jit(jax.vmap(smooth.kinematics, in_axes=(None, 0)))( @@ -224,7 +224,7 @@ class SmoothTest(parameterized.TestCase): d.mocap_pos[:] = dx.mocap_pos d.mocap_quat[:] = dx.mocap_quat m.geom_pos[:] = mx.geom_pos[i] - mujoco.mj_forward(m, d) + mujoco.mj_kinematics(m, d) tu.assert_attr_eq(d, dx, 'xanchor') tu.assert_attr_eq(d, dx, 'xaxis') diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index f13e6df2..df02a52b 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -62,6 +62,7 @@ class BlockDim: contact_sort: int energy_vel_kinetic: int euler_dense: int + linesearch_iterative: int mul_m_dense: int ray: int segmented_sort: int @@ -69,6 +70,7 @@ class BlockDim: update_gradient_JTDAJ_dense: int update_gradient_JTDAJ_sparse: int update_gradient_cholesky: int + update_gradient_cholesky_blocked: int def tree_flatten(self): children = list((getattr(self, k) for k in self.__dataclass_fields__)) @@ -93,6 +95,7 @@ class OptionWarp(PyTreeNode): contact_sensor_maxmatch: int graph_conditional: bool has_fluid: bool + impratio_invsqrt: jax.Array is_sparse: bool ls_parallel: bool ls_parallel_min_step: float @@ -113,6 +116,7 @@ class ModelWarp(PyTreeNode): dof_tri_col: np.ndarray dof_tri_row: np.ndarray eq_connect_adr: np.ndarray + eq_flex_adr: np.ndarray eq_jnt_adr: np.ndarray eq_ten_adr: np.ndarray eq_wld_adr: np.ndarray @@ -122,13 +126,17 @@ class ModelWarp(PyTreeNode): flex_edge: np.ndarray flex_edgeadr: np.ndarray flex_edgeflap: np.ndarray + flex_edgenum: np.ndarray flex_elem: np.ndarray + flex_elemadr: np.ndarray flex_elemedge: np.ndarray flex_elemedgeadr: np.ndarray + flex_elemnum: np.ndarray flex_stiffness: np.ndarray flex_vertadr: np.ndarray flex_vertbodyid: np.ndarray flex_vertnum: np.ndarray + flexedge_invweight0: np.ndarray flexedge_length0: np.ndarray geom_pair_type_count: Tuple[int, ...] geom_plugin_index: np.ndarray @@ -167,6 +175,7 @@ class ModelWarp(PyTreeNode): nsensorcollision: int nsensorcontact: int nsensortaxel: int + nv_pad: int nxn_geom_pair: np.ndarray nxn_geom_pair_filtered: np.ndarray nxn_pairid: np.ndarray @@ -249,8 +258,6 @@ class DataWarp(PyTreeNode): efc__alpha: jax.Array efc__aref: jax.Array efc__beta: jax.Array - efc__cholesky_L_tmp: jax.Array - efc__cholesky_y_tmp: jax.Array efc__cost: jax.Array efc__done: jax.Array efc__force: jax.Array @@ -258,7 +265,6 @@ class DataWarp(PyTreeNode): efc__gauss: jax.Array efc__grad: jax.Array efc__grad_dot: jax.Array - efc__h: jax.Array efc__id: jax.Array efc__jv: jax.Array efc__margin: jax.Array @@ -275,6 +281,7 @@ class DataWarp(PyTreeNode): efc__type: jax.Array efc__vel: jax.Array energy: jax.Array + flexedge_J: jax.Array flexedge_length: jax.Array flexedge_velocity: jax.Array flexvert_xpos: jax.Array @@ -285,6 +292,7 @@ class DataWarp(PyTreeNode): ncollision: jax.Array ne: jax.Array ne_connect: jax.Array + ne_flex: jax.Array ne_jnt: jax.Array ne_ten: jax.Array ne_weld: jax.Array @@ -404,8 +412,6 @@ _NDIM = { 'efc__alpha': 1, 'efc__aref': 2, 'efc__beta': 1, - 'efc__cholesky_L_tmp': 3, - 'efc__cholesky_y_tmp': 2, 'efc__cost': 1, 'efc__done': 1, 'efc__force': 2, @@ -413,7 +419,6 @@ _NDIM = { 'efc__gauss': 1, 'efc__grad': 2, 'efc__grad_dot': 1, - 'efc__h': 3, 'efc__id': 2, 'efc__jv': 2, 'efc__margin': 2, @@ -431,6 +436,7 @@ _NDIM = { 'efc__vel': 2, 'energy': 2, 'eq_active': 2, + 'flexedge_J': 3, 'flexedge_length': 2, 'flexedge_velocity': 2, 'flexvert_xpos': 3, @@ -445,6 +451,7 @@ _NDIM = { 'ncollision': 1, 'ne': 1, 'ne_connect': 1, + 'ne_flex': 1, 'ne_jnt': 1, 'ne_ten': 1, 'ne_weld': 1, @@ -531,6 +538,7 @@ _NDIM = { 'block_dim__contact_sort': 0, 'block_dim__energy_vel_kinetic': 0, 'block_dim__euler_dense': 0, + 'block_dim__linesearch_iterative': 0, 'block_dim__mul_m_dense': 0, 'block_dim__ray': 0, 'block_dim__segmented_sort': 0, @@ -538,6 +546,7 @@ _NDIM = { 'block_dim__update_gradient_JTDAJ_dense': 0, 'block_dim__update_gradient_JTDAJ_sparse': 0, 'block_dim__update_gradient_cholesky': 0, + 'block_dim__update_gradient_cholesky_blocked': 0, 'body_conaffinity': 1, 'body_contype': 1, 'body_dofadr': 1, @@ -589,6 +598,7 @@ _NDIM = { 'eq_active0': 1, 'eq_connect_adr': 1, 'eq_data': 3, + 'eq_flex_adr': 1, 'eq_jnt_adr': 1, 'eq_obj1id': 1, 'eq_obj2id': 1, @@ -605,13 +615,17 @@ _NDIM = { 'flex_edge': 2, 'flex_edgeadr': 1, 'flex_edgeflap': 2, + 'flex_edgenum': 1, 'flex_elem': 1, + 'flex_elemadr': 1, 'flex_elemedge': 1, 'flex_elemedgeadr': 1, + 'flex_elemnum': 1, 'flex_stiffness': 2, 'flex_vertadr': 1, 'flex_vertbodyid': 1, 'flex_vertnum': 1, + 'flexedge_invweight0': 1, 'flexedge_length0': 1, 'geom_aabb': 4, 'geom_bodyid': 1, @@ -743,6 +757,7 @@ _NDIM = { 'ntendon': 0, 'nu': 0, 'nv': 0, + 'nv_pad': 0, 'nwrap': 0, 'nxn_geom_pair': 2, 'nxn_geom_pair_filtered': 2, @@ -761,7 +776,7 @@ _NDIM = { 'opt__graph_conditional': 0, 'opt__gravity': 2, 'opt__has_fluid': 0, - 'opt__impratio': 1, + 'opt__impratio_invsqrt': 1, 'opt__integrator': 0, 'opt__is_sparse': 0, 'opt__iterations': 0, @@ -878,7 +893,7 @@ _NDIM = { 'graph_conditional': 0, 'gravity': 2, 'has_fluid': 0, - 'impratio': 1, + 'impratio_invsqrt': 1, 'integrator': 0, 'is_sparse': 0, 'iterations': 0, @@ -942,8 +957,6 @@ _BATCH_DIM = { 'efc__alpha': True, 'efc__aref': True, 'efc__beta': True, - 'efc__cholesky_L_tmp': True, - 'efc__cholesky_y_tmp': True, 'efc__cost': True, 'efc__done': True, 'efc__force': True, @@ -951,7 +964,6 @@ _BATCH_DIM = { 'efc__gauss': True, 'efc__grad': True, 'efc__grad_dot': True, - 'efc__h': True, 'efc__id': True, 'efc__jv': True, 'efc__margin': True, @@ -969,6 +981,7 @@ _BATCH_DIM = { 'efc__vel': True, 'energy': True, 'eq_active': True, + 'flexedge_J': True, 'flexedge_length': True, 'flexedge_velocity': True, 'flexvert_xpos': True, @@ -983,6 +996,7 @@ _BATCH_DIM = { 'ncollision': False, 'ne': True, 'ne_connect': True, + 'ne_flex': True, 'ne_jnt': True, 'ne_ten': True, 'ne_weld': True, @@ -1069,6 +1083,7 @@ _BATCH_DIM = { 'block_dim__contact_sort': False, 'block_dim__energy_vel_kinetic': False, 'block_dim__euler_dense': False, + 'block_dim__linesearch_iterative': False, 'block_dim__mul_m_dense': False, 'block_dim__ray': False, 'block_dim__segmented_sort': False, @@ -1076,6 +1091,7 @@ _BATCH_DIM = { 'block_dim__update_gradient_JTDAJ_dense': False, 'block_dim__update_gradient_JTDAJ_sparse': False, 'block_dim__update_gradient_cholesky': False, + 'block_dim__update_gradient_cholesky_blocked': False, 'body_conaffinity': False, 'body_contype': False, 'body_dofadr': False, @@ -1127,6 +1143,7 @@ _BATCH_DIM = { 'eq_active0': False, 'eq_connect_adr': False, 'eq_data': True, + 'eq_flex_adr': False, 'eq_jnt_adr': False, 'eq_obj1id': False, 'eq_obj2id': False, @@ -1143,13 +1160,17 @@ _BATCH_DIM = { 'flex_edge': False, 'flex_edgeadr': False, 'flex_edgeflap': False, + 'flex_edgenum': False, 'flex_elem': False, + 'flex_elemadr': False, 'flex_elemedge': False, 'flex_elemedgeadr': False, + 'flex_elemnum': False, 'flex_stiffness': False, 'flex_vertadr': False, 'flex_vertbodyid': False, 'flex_vertnum': False, + 'flexedge_invweight0': False, 'flexedge_length0': False, 'geom_aabb': True, 'geom_bodyid': False, @@ -1281,6 +1302,7 @@ _BATCH_DIM = { 'ntendon': False, 'nu': False, 'nv': False, + 'nv_pad': False, 'nwrap': False, 'nxn_geom_pair': False, 'nxn_geom_pair_filtered': False, @@ -1299,7 +1321,7 @@ _BATCH_DIM = { 'opt__graph_conditional': False, 'opt__gravity': True, 'opt__has_fluid': False, - 'opt__impratio': True, + 'opt__impratio_invsqrt': True, 'opt__integrator': False, 'opt__is_sparse': False, 'opt__iterations': False, @@ -1416,7 +1438,7 @@ _BATCH_DIM = { 'graph_conditional': False, 'gravity': True, 'has_fluid': False, - 'impratio': True, + 'impratio_invsqrt': True, 'integrator': False, 'is_sparse': False, 'iterations': False,