chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Import google-deepmind/mujoco_warp from GitHub.

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