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

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 835230315
Change-Id: Ifa4f8f66c4e3de1a3074f0897daf93f88e469f06
This commit is contained in:
Taylor Howell
2025-11-21 08:15:42 -08:00
committed by Copybara-Service
parent bb488fec29
commit a359780ca1
20 changed files with 1578 additions and 2004 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+284 -62
View File
@@ -32,17 +32,18 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
# TODO(team): improve compile time to enable backward pass
wp.set_module_options({"enable_backward": False})
MULTI_CONTACT_COUNT = 8
mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float)
mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
_CONVEX_COLLISION_PAIRS = [
(GeomType.HFIELD, GeomType.SPHERE),
@@ -260,9 +261,9 @@ def ccd_kernel_builder(
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
) -> int:
points = mat3c()
witness1 = mat3c()
witness2 = mat3c()
points = mat43()
witness1 = mat43()
witness2 = mat43()
geom1.margin = margin
geom2.margin = margin
if pairid[1] >= 0:
@@ -566,6 +567,33 @@ def ccd_kernel_builder(
adr = hfield_adr[geom1_dataid]
hfield_contact_dist = vec_maxconpair()
hfield_contact_pos = mat_maxconpair()
hfield_contact_normal = mat_maxconpair()
min_dist = float(wp.inf)
min_normal = wp.vec3(wp.inf, wp.inf, wp.inf)
min_pos = wp.vec3(wp.inf, wp.inf, wp.inf)
min_id = int(-1)
# TODO(team): height field margin?
geom1.margin = margin
geom2.margin = margin
# EPA memory
epa_vert = epa_vert_in[tid]
epa_vert1 = epa_vert1_in[tid]
epa_vert2 = epa_vert2_in[tid]
epa_vert_index1 = epa_vert_index1_in[tid]
epa_vert_index2 = epa_vert_index2_in[tid]
epa_face = epa_face_in[tid]
epa_pr = epa_pr_in[tid]
epa_norm2 = epa_norm2_in[tid]
epa_index = epa_index_in[tid]
epa_map = epa_map_in[tid]
epa_horizon = epa_horizon_in[tid]
collision_pairid = collision_pairid_in[tid]
# process all prisms in subgrid
count = int(0)
for r in range(rmin, rmax):
@@ -573,6 +601,13 @@ def ccd_kernel_builder(
for c in range(cmin, cmax + 1):
# add both triangles from this cell
for i in range(2):
if count >= MJ_MAXCONPAIR:
wp.printf(
"height field collision overflow, number of collisions >= %u - please adjust resolution: \n decrease the number of hfield rows/cols or modify size of colliding geom\n",
MJ_MAXCONPAIR,
)
continue
# add vert
x = dx * float(c) - size[0]
y = dy * float(r + dr[i]) - size[1]
@@ -607,66 +642,253 @@ def ccd_kernel_builder(
x1_ += prism[i]
x1 += geom1.rot @ (x1_ / 6.0)
ncontact = eval_ccd_write_contact(
opt_ccd_tolerance,
geom_type,
naconmax_in,
epa_vert_in,
epa_vert1_in,
epa_vert2_in,
epa_vert_index1_in,
epa_vert_index2_in,
epa_face_in,
epa_pr_in,
epa_norm2_in,
epa_index_in,
epa_map_in,
epa_horizon_in,
multiccd_polygon_in,
multiccd_clipped_in,
multiccd_pnormal_in,
multiccd_pdist_in,
multiccd_idx1_in,
multiccd_idx2_in,
multiccd_n1_in,
multiccd_n2_in,
multiccd_endvert_in,
multiccd_face1_in,
multiccd_face2_in,
dist, ncontact, w1, w2, idx = ccd(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
0.0,
ccd_iterations,
geom1,
geom2,
geoms,
worldid,
tid,
margin,
gap,
condim,
friction,
solref,
solreffriction,
solimp,
geomtype1,
geomtype2,
x1,
geom2.pos,
count,
collision_pairid_in[tid],
contact_dist_out,
contact_pos_out,
contact_frame_out,
contact_includemargin_out,
contact_friction_out,
contact_solref_out,
contact_solreffriction_out,
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
epa_vert,
epa_vert1,
epa_vert2,
epa_vert_index1,
epa_vert_index2,
epa_face,
epa_pr,
epa_norm2,
epa_index,
epa_map,
epa_horizon,
)
count += ncontact
if count >= MJ_MAXCONPAIR:
return
if ncontact == 0:
continue
# cache contact information
hfield_contact_dist[count] = dist
pos = 0.5 * (w1 + w2)
hfield_contact_pos[count, 0] = pos[0]
hfield_contact_pos[count, 1] = pos[1]
hfield_contact_pos[count, 2] = pos[2]
frame = make_frame(w1 - w2)
normal = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2])
hfield_contact_normal[count, 0] = normal[0]
hfield_contact_normal[count, 1] = normal[1]
hfield_contact_normal[count, 2] = normal[2]
# contact with minimum distance
if dist < min_dist:
min_dist = dist
min_normal = normal
min_pos = pos
min_id = count
count += 1
# contact 0: minimum distance
write_contact(
naconmax_in,
0,
min_dist,
min_pos,
make_frame(min_normal),
margin,
gap,
condim,
friction,
solref,
solreffriction,
solimp,
geoms,
collision_pairid,
worldid,
contact_dist_out,
contact_pos_out,
contact_frame_out,
contact_includemargin_out,
contact_friction_out,
contact_solref_out,
contact_solreffriction_out,
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
# TODO(team): routine for select subset of contacts
# TODO(team): if use_multiccd?
if wp.static(True):
MIN_DIST_TO_NEXT_CONTACT = 1.0e-3
# contact 1: furthest from minimum distance contact
id1 = int(-1)
dist1 = float(-wp.inf)
for i in range(count):
if i == min_id:
continue
hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2])
dist = wp.norm_l2(hf_pos - min_pos)
if dist > dist1:
id1 = i
dist1 = dist
if id1 == -1 or (0.0 < dist1 and dist1 < MIN_DIST_TO_NEXT_CONTACT):
return
pos1 = wp.vec3(hfield_contact_pos[id1, 0], hfield_contact_pos[id1, 1], hfield_contact_pos[id1, 2])
normal1 = wp.vec3(hfield_contact_normal[id1, 0], hfield_contact_normal[id1, 1], hfield_contact_normal[id1, 2])
write_contact(
naconmax_in,
1,
hfield_contact_dist[id1],
pos1,
make_frame(normal1),
margin,
gap,
condim,
friction,
solref,
solreffriction,
solimp,
geoms,
collision_pairid,
worldid,
contact_dist_out,
contact_pos_out,
contact_frame_out,
contact_includemargin_out,
contact_friction_out,
contact_solref_out,
contact_solreffriction_out,
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
# contact 2: point furthest from min_pos - pos1 line
dist_min1 = wp.cross(min_normal, min_pos - pos1)
id2 = int(-1)
dist_12 = float(-wp.inf)
for i in range(count):
if i == min_id or i == id1:
continue
hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2])
dist = wp.abs(wp.dot(hf_pos - min_pos, dist_min1))
if dist > dist_12:
id2 = i
dist_12 = dist
if id2 == -1 or (0.0 < dist_12 and dist_12 < MIN_DIST_TO_NEXT_CONTACT):
return
pos2 = wp.vec3(hfield_contact_pos[id2, 0], hfield_contact_pos[id2, 1], hfield_contact_pos[id2, 2])
normal2 = wp.vec3(hfield_contact_normal[id2, 0], hfield_contact_normal[id2, 1], hfield_contact_normal[id2, 2])
write_contact(
naconmax_in,
2,
hfield_contact_dist[id2],
pos2,
make_frame(normal2),
margin,
gap,
condim,
friction,
solref,
solreffriction,
solimp,
geoms,
collision_pairid,
worldid,
contact_dist_out,
contact_pos_out,
contact_frame_out,
contact_includemargin_out,
contact_friction_out,
contact_solref_out,
contact_solreffriction_out,
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
# contact 3: point furthest from other triangle edge
vec_min2 = wp.cross(min_normal, min_pos - pos2)
vec_12 = wp.cross(min_normal, pos1 - pos2)
id3 = int(-1)
dist3 = float(-wp.inf)
for i in range(count):
if i == min_id or i == id1 or i == id2:
continue
hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2])
dist = wp.abs(wp.dot(hf_pos - min_pos, vec_min2)) + wp.abs(wp.dot(pos1 - hf_pos, vec_12))
if dist > dist3:
id3 = i
dist3 = dist
if id3 == -1 or (0.0 < dist3 and dist3 < MIN_DIST_TO_NEXT_CONTACT):
return
pos3 = wp.vec3(hfield_contact_pos[id3, 0], hfield_contact_pos[id3, 1], hfield_contact_pos[id3, 2])
normal3 = wp.vec3(hfield_contact_normal[id3, 0], hfield_contact_normal[id3, 1], hfield_contact_normal[id3, 2])
write_contact(
naconmax_in,
3,
hfield_contact_dist[id3],
pos3,
make_frame(normal3),
margin,
gap,
condim,
friction,
solref,
solreffriction,
solimp,
geoms,
collision_pairid,
worldid,
contact_dist_out,
contact_pos_out,
contact_frame_out,
contact_includemargin_out,
contact_friction_out,
contact_solref_out,
contact_solreffriction_out,
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
else:
eval_ccd_write_contact(
opt_ccd_tolerance,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+21 -21
View File
@@ -27,9 +27,11 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -170,10 +172,6 @@ def _aabb_filter(
return True
mat23 = wp.types.matrix(shape=(2, 3), dtype=float)
mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
# TODO(team): improve performance by precomputing bounding box
@wp.func
def _obb_filter(
@@ -232,7 +230,7 @@ def _obb_filter(
return True
def _broadphase_filter(m: Model):
def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
@wp.func
def func(
# Model:
@@ -252,28 +250,28 @@ def _broadphase_filter(m: Model):
# 4: aabb
# 8: obb
aabb_id = worldid % geom_aabb.shape[0] if wp.static(m.geom_aabb.shape[0] > 1) else 0
aabb_id = worldid % ngeom_aabb if wp.static(ngeom_aabb > 1) else 0
center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0]
size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1]
rbound_id = worldid % geom_rbound.shape[0] if wp.static(m.geom_rbound.shape[0] > 1) else 0
rbound_id = worldid % ngeom_rbound if wp.static(ngeom_rbound > 1) else 0
rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2]
margin_id = worldid % geom_margin.shape[0] if wp.static(m.geom_margin.shape[0] > 1) else 0
margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2]
xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
if rbound1 == 0.0 or rbound2 == 0.0:
if wp.static(m.opt.broadphase_filter & BroadphaseFilter.PLANE):
if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE):
return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
else:
if wp.static(m.opt.broadphase_filter & BroadphaseFilter.SPHERE):
if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE):
if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
return False
if wp.static(m.opt.broadphase_filter & BroadphaseFilter.AABB):
if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB):
if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
if wp.static(m.opt.broadphase_filter & BroadphaseFilter.OBB):
if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB):
if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
@@ -403,7 +401,7 @@ def _sap_range(
@cache_kernel
def _sap_broadphase(broadphase_filter):
def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
@nested_kernel(module="unique", enable_backward=False)
def kernel(
# Model:
@@ -461,7 +459,9 @@ def _sap_broadphase(broadphase_filter):
continue
if (
broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid)
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
)
or pairid[1] >= 0
):
_add_geom_pair(
@@ -578,9 +578,8 @@ def sap_broadphase(m: Model, d: Data):
# estimate number of overlap checks
# assumes each geom has 5 other geoms (batched over all worlds)
nsweep = 5 * nworldgeom
broadphase_filter = _broadphase_filter(m)
wp.launch(
kernel=_sap_broadphase(broadphase_filter),
kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
dim=nsweep,
inputs=[
m.ngeom,
@@ -602,7 +601,7 @@ def sap_broadphase(m: Model, d: Data):
@cache_kernel
def _nxn_broadphase(broadphase_filter):
def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
@nested_kernel(module="unique", enable_backward=False)
def kernel(
# Model:
@@ -629,7 +628,9 @@ def _nxn_broadphase(broadphase_filter):
geom2 = geom[1]
if (
broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid)
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
)
or nxn_pairid[elementid][1] >= 0
):
_add_geom_pair(
@@ -663,9 +664,8 @@ def nxn_broadphase(m: Model, d: Data):
The initial list of pairs is filtered at model creation time to exclude pairs based on
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
"""
broadphase_filter = _broadphase_filter(m)
wp.launch(
_nxn_broadphase(broadphase_filter),
_nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[
m.geom_type,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+22 -32
View File
@@ -20,6 +20,8 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
# TODO(team): improve compile time to enable backward pass
wp.set_module_options({"enable_backward": False})
@@ -32,13 +34,6 @@ MJ_MINVAL2 = MJ_MINVAL * MJ_MINVAL
FACE_TOL = 0.99999872
EDGE_TOL = 0.00159999931
mat43 = wp.types.matrix(shape=(4, 3), dtype=float)
mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
# box-box currently supports up to 8 contacts
MULTI_CONTACT_COUNT = 8
mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
@wp.struct
class GJKResult:
@@ -155,30 +150,25 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
vert_edgeadr = geom.graphadr + 2
vert_globalid = geom.graphadr + 2 + numvert
edge_localid = geom.graphadr + 2 + 2 * numvert
# hillclimb until no change
prev = int(-1)
imax = int(0)
if geom.index > -1:
imax = geom.index
sp.cached_index = geom.index
imax = wp.where(geom.index > -1, geom.index, 0)
while True:
prev = int(imax)
i = int(geom.graph[vert_edgeadr + imax])
while geom.graph[edge_localid + i] >= 0:
subidx = geom.graph[edge_localid + i]
# hillclimb until no change
while imax != prev:
prev = imax
i = geom.graph[vert_edgeadr + imax]
subidx = geom.graph[edge_localid + i]
while subidx >= 0:
idx = geom.graph[vert_globalid + subidx]
dist = wp.dot(local_dir, geom.vert[geom.vertadr + idx])
if dist > max_dist:
max_dist = dist
imax = int(subidx)
i += int(1)
if imax == prev:
break
imax = wp.where(dist > max_dist, subidx, imax)
max_dist = wp.where(dist > max_dist, dist, max_dist)
i += 1
subidx = geom.graph[edge_localid + i]
sp.cached_index = imax
imax = geom.graph[vert_globalid + imax]
sp.vertex_index = imax
sp.point = geom.vert[geom.vertadr + imax]
sp.vertex_index = geom.graph[vert_globalid + imax]
sp.point = geom.vert[geom.vertadr + sp.vertex_index]
sp.point = geom.rot @ sp.point + geom.pos
elif geomtype == GeomType.HFIELD:
@@ -1813,9 +1803,9 @@ def _polygon_clip(
# Out:
polygon_out: wp.array(dtype=wp.vec3),
clipped_out: wp.array(dtype=wp.vec3),
) -> Tuple[int, mat3c, mat3c]:
witness1 = mat3c()
witness2 = mat3c()
) -> Tuple[int, mat43, mat43]:
witness1 = mat43()
witness2 = mat43()
# clipping face needs to be at least a triangle
if nface1 < 3:
@@ -1930,9 +1920,9 @@ def multicontact(
geom2: Geom,
geomtype1: int,
geomtype2: int,
) -> Tuple[int, mat3c, mat3c]:
witness1 = mat3c()
witness2 = mat3c()
) -> Tuple[int, mat43, mat43]:
witness1 = mat43()
witness2 = mat43()
witness1[0] = x1
witness2[0] = x2
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+5 -10
View File
@@ -38,21 +38,16 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
class mat43f(wp.types.matrix(shape=(4, 3), dtype=wp.float32)):
pass
mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
@wp.struct
class Geom:
pos: wp.vec3
@@ -169,7 +164,7 @@ def geom_collision_pair(
@wp.func
def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43f, wp.vec3]:
def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43, wp.vec3]:
"""Core contact geometry calculation for plane-convex collision.
Args:
@@ -185,7 +180,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
_HUGE_VAL = 1e6
contact_dist = wp.vec4(wp.inf)
contact_pos = mat43f()
contact_pos = mat43()
contact_count = int(0)
# get points in the convex frame
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+8 -8
View File
@@ -26,7 +26,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8f
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -54,7 +54,7 @@ class VolumeData:
half_size: wp.vec3
oct_aabb: wp.array2d(dtype=wp.vec3)
oct_child: wp.array(dtype=vec8i)
oct_coeff: wp.array(dtype=vec8f)
oct_coeff: wp.array(dtype=vec8)
valid: bool = False
@@ -79,7 +79,7 @@ def get_sdf_params(
# Model:
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8f),
oct_coeff: wp.array(dtype=vec8),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
# In:
@@ -225,12 +225,12 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
@wp.func
def find_oct(
oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool
) -> Tuple[int, Tuple[vec8f, vec8f, vec8f]]:
) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
stack = int(0)
niter = int(100)
rx = vec8f(0.0)
ry = vec8f(0.0)
rz = vec8f(0.0)
rx = vec8(0.0)
ry = vec8(0.0)
rz = vec8(0.0)
eps = 1e-6
while niter > 0:
@@ -621,7 +621,7 @@ def _sdf_narrowphase(
nmeshface: int,
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8f),
oct_coeff: wp.array(dtype=vec8),
geom_type: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+2 -2
View File
@@ -1843,7 +1843,7 @@ def make_constraint(m: types.Model, d: types.Data):
if m.opt.cone == types.ConeType.PYRAMIDAL:
wp.launch(
_efc_contact_pyramidal,
dim=(d.naconmax, 2 * (m.condim_max - 1) if m.condim_max > 1 else 1),
dim=(d.naconmax, m.nmaxpyramid),
inputs=[
m.nv,
m.opt.timestep,
@@ -1888,7 +1888,7 @@ def make_constraint(m: types.Model, d: types.Data):
elif m.opt.cone == types.ConeType.ELLIPTIC:
wp.launch(
_efc_contact_elliptic,
dim=(d.naconmax, m.condim_max),
dim=(d.naconmax, m.nmaxcondim),
inputs=[
m.nv,
m.opt.timestep,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+14 -7
View File
@@ -75,14 +75,21 @@ def _qderiv_actuator_passive(
else:
bias = 0.0
if actuator_dyntype[actid] != DynType.NONE:
act_first = actuator_actadr[actid]
act_last = act_first + actuator_actnum[actid] - 1
vel = bias + gain * act_in[worldid, act_last]
else:
vel = bias + gain * ctrl_in[worldid, actid]
if bias == 0.0 and gain == 0.0:
continue
qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel
vel = bias
if actuator_dyntype[actid] != DynType.NONE:
if gain != 0.0:
act_first = actuator_actadr[actid]
act_last = act_first + actuator_actnum[actid] - 1
vel += gain * act_in[worldid, act_last]
else:
if gain != 0.0:
vel += gain * ctrl_in[worldid, actid]
if vel != 0.0:
qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel
# TODO(team): fluid model derivative
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+26 -25
View File
@@ -42,7 +42,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -521,7 +521,8 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
# TODO(team): sparse actuator_moment version
def _actuator_velocity(m: Model, d: Data):
@cache_kernel
def _actuator_velocity(nv: int):
@nested_kernel(module="unique", enable_backward=False)
def actuator_velocity(
# Data in:
@@ -531,22 +532,17 @@ def _actuator_velocity(m: Model, d: Data):
actuator_velocity_out: wp.array2d(dtype=float),
):
worldid, actid = wp.tid()
moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(m.nv))
qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(m.nv))
moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(nv))
qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
moment_qvel_tile = wp.tile_map(wp.mul, moment_tile, qvel_tile)
actuator_velocity_tile = wp.tile_reduce(wp.add, moment_qvel_tile)
actuator_velocity_out[worldid, actid] = actuator_velocity_tile[0]
wp.launch_tiled(
actuator_velocity,
dim=(d.nworld, m.nu),
inputs=[d.qvel, d.actuator_moment],
outputs=[d.actuator_velocity],
block_dim=m.block_dim.actuator_velocity,
)
return actuator_velocity
def _tendon_velocity(m: Model, d: Data):
@cache_kernel
def _tendon_velocity(nv: int):
@nested_kernel(module="unique", enable_backward=False)
def tendon_velocity(
# Data in:
@@ -556,29 +552,34 @@ def _tendon_velocity(m: Model, d: Data):
ten_velocity_out: wp.array2d(dtype=float),
):
worldid, tenid = wp.tid()
ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(m.nv))
qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(m.nv))
ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(nv))
qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
ten_J_qvel_tile = wp.tile_map(wp.mul, ten_J_tile, qvel_tile)
ten_velocity_tile = wp.tile_reduce(wp.add, ten_J_qvel_tile)
ten_velocity_out[worldid, tenid] = ten_velocity_tile[0]
wp.launch_tiled(
tendon_velocity,
dim=(d.nworld, m.ntendon),
inputs=[d.qvel, d.ten_J],
outputs=[d.ten_velocity],
block_dim=m.block_dim.tendon_velocity,
)
return tendon_velocity
@event_scope
def fwd_velocity(m: Model, d: Data):
"""Velocity-dependent computations."""
_actuator_velocity(m, d)
wp.launch_tiled(
_actuator_velocity(m.nv),
dim=(d.nworld, m.nu),
inputs=[d.qvel, d.actuator_moment],
outputs=[d.actuator_velocity],
block_dim=m.block_dim.actuator_velocity,
)
if m.ntendon > 0:
# TODO(team): sparse version
_tendon_velocity(m, d)
# TODO(team): sparse version
wp.launch_tiled(
_tendon_velocity(m.nv),
dim=(d.nworld, m.ntendon),
inputs=[d.qvel, d.ten_J],
outputs=[d.ten_velocity],
block_dim=m.block_dim.tendon_velocity,
)
smooth.com_vel(m, d)
passive.passive(m, d)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+527 -1223
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+7 -7
View File
@@ -561,7 +561,7 @@ def _flex_elasticity(
flex_elem: wp.array(dtype=int),
flex_elemedge: wp.array(dtype=int),
flexedge_length0: wp.array(dtype=float),
flex_stiffness: wp.array(dtype=float),
flex_stiffness: wp.array2d(dtype=float),
flex_damping: wp.array(dtype=float),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -612,8 +612,8 @@ def _flex_elasticity(
id = int(0)
for ed1 in range(nedge):
for ed2 in range(ed1, nedge):
metric[ed1, ed2] = flex_stiffness[21 * elemid + id]
metric[ed2, ed1] = flex_stiffness[21 * elemid + id]
metric[ed1, ed2] = flex_stiffness[elemid, id]
metric[ed2, ed1] = flex_stiffness[elemid, id]
id += 1
force = wp.mat(0.0, shape=(6, 3))
@@ -640,7 +640,7 @@ def _flex_bending(
flex_vertbodyid: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_edgeflap: wp.array(dtype=wp.vec2i),
flex_bending: wp.array(dtype=float),
flex_bending: wp.array2d(dtype=float),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# Data out:
@@ -664,7 +664,7 @@ def _flex_bending(
return
frc = wp.mat(0.0, shape=(4, 3))
if flex_bending[17 * edgeid + 16]:
if flex_bending[edgeid, 16]:
v0 = flexvert_xpos_in[worldid, v[0]]
v1 = flexvert_xpos_in[worldid, v[1]]
v2 = flexvert_xpos_in[worldid, v[2]]
@@ -678,8 +678,8 @@ def _flex_bending(
for i in range(nvert):
for x in range(3):
for j in range(nvert):
force[i, x] -= flex_bending[17 * edgeid + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] -= flex_bending[17 * edgeid + 16] * frc[i, x]
force[i, x] -= flex_bending[edgeid, 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] -= flex_bending[edgeid, 16] * frc[i, x]
for i in range(nvert):
bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+5 -5
View File
@@ -37,12 +37,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType
from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8f
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -482,8 +482,8 @@ def _sensor_pos(
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_pos_adr: wp.array(dtype=int),
sensor_collision_start_adr: wp.array(dtype=int),
rangefinder_sensor_adr: wp.array(dtype=int),
sensor_collision_start_adr: wp.array(dtype=int),
collision_sensor_adr: wp.array(dtype=int),
# Data in:
time_in: wp.array(dtype=float),
@@ -844,8 +844,8 @@ def sensor_pos(m: Model, d: Data):
m.sensor_adr,
m.sensor_cutoff,
m.sensor_pos_adr,
m.sensor_collision_start_adr,
m.rangefinder_sensor_adr,
m.sensor_collision_start_adr,
m.collision_sensor_adr,
d.time,
d.energy,
@@ -2089,7 +2089,7 @@ def _sensor_tactile(
body_weldid: wp.array(dtype=int),
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8f),
oct_coeff: wp.array(dtype=vec8),
geom_type: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+1 -1
View File
@@ -35,7 +35,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.types import vec11
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+7 -6
View File
@@ -27,7 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocke
from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -1448,7 +1448,7 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
@cache_kernel
def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int):
if njmax < tile_size:
tile_size = njmax
@@ -1473,7 +1473,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
nefc = nefc_in[worldid]
sum_val = wp.tile_load(qM_in[worldid], shape=(nv, nv), bounds_check=False)
sum_val = wp.tile_load(qM_in[worldid], shape=(nv_padded, nv_padded), bounds_check=True)
# Each tile processes one output tile by looping over all constraints
for k in range(0, njmax, TILE_SIZE_K):
@@ -1483,7 +1483,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
# AD: leaving bounds-check disabled here because I'm not entirely sure that
# everything always hits the fast path. The padding takes care of any
# potential OOB accesses.
J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv), offset=(k, 0), bounds_check=False)
J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv_padded), offset=(k, 0), bounds_check=False)
J_kj = J_ki
# state check
@@ -1499,7 +1499,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
D_k = wp.tile_map(wp.mul, active_tile, D_k)
J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(nv, TILE_SIZE_K)))
J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(nv_padded, TILE_SIZE_K)))
sum_val += wp.tile_matmul(J_ki, J_kj)
@@ -1742,8 +1742,9 @@ def _update_gradient(m: types.Model, d: types.Data):
outputs=[d.efc.h],
)
else:
nv_padded = d.efc.J.shape[2]
wp.launch_tiled(
update_gradient_JTDAJ_dense_tiled(m.nv, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
update_gradient_JTDAJ_dense_tiled(nv_padded, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
dim=d.nworld,
inputs=[
d.nefc,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+1 -1
View File
@@ -27,7 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+598 -553
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
Vendored
+6 -6
View File
@@ -121,9 +121,9 @@ def event_scope(fn, name: str = ""):
return wrapper
# @kernel decorator to automatically set up modules based on nested
# @nested_kernel decorator to automatically set up modules based on nested
# function names
def kernel(
def nested_kernel(
f: Optional[Callable] = None,
*,
enable_backward: Optional[bool] = None,
@@ -136,20 +136,20 @@ def kernel(
Example::
@kernel
@nested_kernel
def my_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
tid = wp.tid()
b[tid] = a[tid] + 1.0
@kernel(enable_backward=False)
@nested_kernel(enable_backward=False)
def my_kernel_no_backward(a: wp.array(dtype=float, ndim=2), x: float):
# the backward pass will not be generated
i, j = wp.tid()
a[i, j] = x
@kernel(module="unique")
@nested_kernel(module="unique")
def my_kernel_unique_module(a: wp.array(dtype=float), b: wp.array(dtype=float)):
# the kernel will be registered in new unique module created just for this
# kernel and its dependent functions and structs
@@ -157,7 +157,7 @@ def kernel(
b[tid] = a[tid] + 1.0
@kernel(enable_backward=False, module=None)
@neste_kernel(enable_backward=False, module=None)
def my_kernel_with_args(a: wp.array(dtype=float), b: wp.array(dtype=float)):
# can now use arguments even when module=None
tid = wp.tid()
mjx/mujoco/mjx/warp/collision_driver.py
+1 -1
View File
@@ -94,7 +94,7 @@ def _collision_shim(
nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_child: wp.array(dtype=mjwp_types.vec8i),
oct_coeff: wp.array(dtype=mjwp_types.vec8f),
oct_coeff: wp.array(dtype=mjwp_types.vec8),
pair_dim: wp.array(dtype=int),
pair_friction: wp.array2d(dtype=mjwp_types.vec5),
pair_gap: wp.array2d(dtype=float),
mjx/mujoco/mjx/warp/forward.py
+18 -12
View File
@@ -106,7 +106,6 @@ def _forward_shim(
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_targetbodyid: wp.array(dtype=int),
collision_sensor_adr: wp.array(dtype=int),
condim_max: int,
dof_Madr: wp.array(dtype=int),
dof_armature: wp.array2d(dtype=float),
dof_bodyid: wp.array(dtype=int),
@@ -129,7 +128,7 @@ def _forward_shim(
eq_solref: wp.array2d(dtype=wp.vec2),
eq_ten_adr: wp.array(dtype=int),
eq_wld_adr: wp.array(dtype=int),
flex_bending: wp.array(dtype=float),
flex_bending: wp.array2d(dtype=float),
flex_damping: wp.array(dtype=float),
flex_dim: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
@@ -138,7 +137,7 @@ def _forward_shim(
flex_elem: wp.array(dtype=int),
flex_elemedge: wp.array(dtype=int),
flex_elemedgeadr: wp.array(dtype=int),
flex_stiffness: wp.array(dtype=float),
flex_stiffness: wp.array2d(dtype=float),
flex_vertadr: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flexedge_length0: wp.array(dtype=float),
@@ -229,8 +228,10 @@ def _forward_shim(
ngravcomp: int,
njnt: int,
nlight: int,
nmaxcondim: int,
nmaxmeshdeg: int,
nmaxpolygon: int,
nmaxpyramid: int,
nmeshface: int,
nmocap: int,
nrangefinder: int,
@@ -247,7 +248,7 @@ def _forward_shim(
nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_child: wp.array(dtype=mjwp_types.vec8i),
oct_coeff: wp.array(dtype=mjwp_types.vec8f),
oct_coeff: wp.array(dtype=mjwp_types.vec8),
pair_dim: wp.array(dtype=int),
pair_friction: wp.array2d(dtype=mjwp_types.vec5),
pair_gap: wp.array2d(dtype=float),
@@ -556,7 +557,6 @@ def _forward_shim(
_m.cam_sensorsize = cam_sensorsize
_m.cam_targetbodyid = cam_targetbodyid
_m.collision_sensor_adr = collision_sensor_adr
_m.condim_max = condim_max
_m.dof_Madr = dof_Madr
_m.dof_armature = dof_armature
_m.dof_bodyid = dof_bodyid
@@ -679,8 +679,10 @@ def _forward_shim(
_m.ngravcomp = ngravcomp
_m.njnt = njnt
_m.nlight = nlight
_m.nmaxcondim = nmaxcondim
_m.nmaxmeshdeg = nmaxmeshdeg
_m.nmaxpolygon = nmaxpolygon
_m.nmaxpyramid = nmaxpyramid
_m.nmeshface = nmeshface
_m.nmocap = nmocap
_m.nrangefinder = nrangefinder
@@ -1279,7 +1281,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.cam_sensorsize,
m.cam_targetbodyid,
m._impl.collision_sensor_adr,
m._impl.condim_max,
m.dof_Madr,
m.dof_armature,
m.dof_bodyid,
@@ -1402,8 +1403,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.ngravcomp,
m.njnt,
m.nlight,
m._impl.nmaxcondim,
m._impl.nmaxmeshdeg,
m._impl.nmaxpolygon,
m._impl.nmaxpyramid,
m.nmeshface,
m.nmocap,
m._impl.nrangefinder,
@@ -1895,7 +1898,6 @@ def _step_shim(
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_targetbodyid: wp.array(dtype=int),
collision_sensor_adr: wp.array(dtype=int),
condim_max: int,
dof_Madr: wp.array(dtype=int),
dof_armature: wp.array2d(dtype=float),
dof_bodyid: wp.array(dtype=int),
@@ -1918,7 +1920,7 @@ def _step_shim(
eq_solref: wp.array2d(dtype=wp.vec2),
eq_ten_adr: wp.array(dtype=int),
eq_wld_adr: wp.array(dtype=int),
flex_bending: wp.array(dtype=float),
flex_bending: wp.array2d(dtype=float),
flex_damping: wp.array(dtype=float),
flex_dim: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
@@ -1927,7 +1929,7 @@ def _step_shim(
flex_elem: wp.array(dtype=int),
flex_elemedge: wp.array(dtype=int),
flex_elemedgeadr: wp.array(dtype=int),
flex_stiffness: wp.array(dtype=float),
flex_stiffness: wp.array2d(dtype=float),
flex_vertadr: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flexedge_length0: wp.array(dtype=float),
@@ -2019,8 +2021,10 @@ def _step_shim(
ngravcomp: int,
njnt: int,
nlight: int,
nmaxcondim: int,
nmaxmeshdeg: int,
nmaxpolygon: int,
nmaxpyramid: int,
nmeshface: int,
nmocap: int,
nrangefinder: int,
@@ -2037,7 +2041,7 @@ def _step_shim(
nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_child: wp.array(dtype=mjwp_types.vec8i),
oct_coeff: wp.array(dtype=mjwp_types.vec8f),
oct_coeff: wp.array(dtype=mjwp_types.vec8),
pair_dim: wp.array(dtype=int),
pair_friction: wp.array2d(dtype=mjwp_types.vec5),
pair_gap: wp.array2d(dtype=float),
@@ -2347,7 +2351,6 @@ def _step_shim(
_m.cam_sensorsize = cam_sensorsize
_m.cam_targetbodyid = cam_targetbodyid
_m.collision_sensor_adr = collision_sensor_adr
_m.condim_max = condim_max
_m.dof_Madr = dof_Madr
_m.dof_armature = dof_armature
_m.dof_bodyid = dof_bodyid
@@ -2471,8 +2474,10 @@ def _step_shim(
_m.ngravcomp = ngravcomp
_m.njnt = njnt
_m.nlight = nlight
_m.nmaxcondim = nmaxcondim
_m.nmaxmeshdeg = nmaxmeshdeg
_m.nmaxpolygon = nmaxpolygon
_m.nmaxpyramid = nmaxpyramid
_m.nmeshface = nmeshface
_m.nmocap = nmocap
_m.nrangefinder = nrangefinder
@@ -3072,7 +3077,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.cam_sensorsize,
m.cam_targetbodyid,
m._impl.collision_sensor_adr,
m._impl.condim_max,
m.dof_Madr,
m.dof_armature,
m.dof_bodyid,
@@ -3196,8 +3200,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.ngravcomp,
m.njnt,
m.nlight,
m._impl.nmaxcondim,
m._impl.nmaxmeshdeg,
m._impl.nmaxpolygon,
m._impl.nmaxpyramid,
m.nmeshface,
m.nmocap,
m._impl.nrangefinder,
mjx/mujoco/mjx/warp/forward_test.py
+2 -1
View File
@@ -155,7 +155,8 @@ class ForwardTest(parameterized.TestCase):
mujoco.mj_fullM(m, qm, d.qM)
else:
qm = d.qM
tu.assert_eq(qm, dx._impl.qM, '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.
tu.assert_contact_eq(d, dx, worldid=i)
mjx/mujoco/mjx/warp/types.py
+23 -21
View File
@@ -105,14 +105,11 @@ class ModelWarp(PyTreeNode):
M_colind: np.ndarray
M_rowadr: np.ndarray
M_rownnz: np.ndarray
actuator_moment_tiles_nu: Tuple[TileSet, ...]
actuator_moment_tiles_nv: Tuple[TileSet, ...]
actuator_trntype_body_adr: np.ndarray
block_dim: BlockDim
body_fluid_ellipsoid: np.ndarray
body_tree: Tuple[np.ndarray, ...]
collision_sensor_adr: np.ndarray
condim_max: int
dof_tri_col: np.ndarray
dof_tri_row: np.ndarray
eq_connect_adr: np.ndarray
@@ -158,9 +155,14 @@ class ModelWarp(PyTreeNode):
nflexedge: int
nflexelem: int
nflexelemdata: int
nflexelemedge: int
nflexvert: int
nmaxcondim: int
nmaxmeshdeg: int
nmaxpolygon: int
nmaxpyramid: int
noct: int
nplugin: int
nrangefinder: int
nsensorcollision: int
nsensorcontact: int
@@ -522,8 +524,6 @@ _NDIM = {
'actuator_gaintype': 1,
'actuator_gear': 3,
'actuator_lengthrange': 2,
'actuator_moment_tiles_nu': -1,
'actuator_moment_tiles_nv': -1,
'actuator_trnid': 2,
'actuator_trntype': 1,
'actuator_trntype_body_adr': 1,
@@ -577,7 +577,6 @@ _NDIM = {
'cam_sensorsize': 2,
'cam_targetbodyid': 1,
'collision_sensor_adr': 1,
'condim_max': 0,
'dof_Madr': 1,
'dof_armature': 2,
'dof_bodyid': 1,
@@ -603,7 +602,7 @@ _NDIM = {
'eq_type': 1,
'eq_wld_adr': 1,
'exclude_signature': 1,
'flex_bending': 1,
'flex_bending': 2,
'flex_damping': 1,
'flex_dim': 1,
'flex_edge': 2,
@@ -612,7 +611,7 @@ _NDIM = {
'flex_elem': 1,
'flex_elemedge': 1,
'flex_elemedgeadr': 1,
'flex_stiffness': 1,
'flex_stiffness': 2,
'flex_vertadr': 1,
'flex_vertbodyid': 1,
'flex_vertnum': 1,
@@ -677,7 +676,7 @@ _NDIM = {
'light_type': 2,
'mapM2M': 1,
'mat_rgba': 3,
'mat_texid': 3,
'mat_texid': 2,
'mat_texrepeat': 3,
'mesh_face': 2,
'mesh_faceadr': 1,
@@ -711,6 +710,7 @@ _NDIM = {
'nflexedge': 0,
'nflexelem': 0,
'nflexelemdata': 0,
'nflexelemedge': 0,
'nflexvert': 0,
'ngeom': 0,
'ngravcomp': 0,
@@ -719,16 +719,22 @@ _NDIM = {
'njnt': 0,
'nlight': 0,
'nmat': 0,
'nmaxcondim': 0,
'nmaxmeshdeg': 0,
'nmaxpolygon': 0,
'nmaxpyramid': 0,
'nmesh': 0,
'nmeshface': 0,
'nmeshgraph': 0,
'nmeshnormal': 0,
'nmeshpoly': 0,
'nmeshpolymap': 0,
'nmeshpolyvert': 0,
'nmeshvert': 0,
'nmocap': 0,
'noct': 0,
'npair': 0,
'nplugin': 0,
'nq': 0,
'nrangefinder': 0,
'nsensor': 0,
@@ -748,8 +754,6 @@ _NDIM = {
'oct_aabb': 3,
'oct_child': 2,
'oct_coeff': 2,
'opt__broadphase': 0,
'opt__broadphase_filter': 0,
'opt__ccd_iterations': 0,
'opt__ccd_tolerance': 1,
'opt__cone': 0,
@@ -867,8 +871,6 @@ _NDIM = {
'wrap_type': 1,
},
'Option': {
'broadphase': 0,
'broadphase_filter': 0,
'ccd_iterations': 0,
'ccd_tolerance': 1,
'cone': 0,
@@ -1060,8 +1062,6 @@ _BATCH_DIM = {
'actuator_gaintype': False,
'actuator_gear': True,
'actuator_lengthrange': False,
'actuator_moment_tiles_nu': False,
'actuator_moment_tiles_nv': False,
'actuator_trnid': False,
'actuator_trntype': False,
'actuator_trntype_body_adr': False,
@@ -1115,7 +1115,6 @@ _BATCH_DIM = {
'cam_sensorsize': False,
'cam_targetbodyid': False,
'collision_sensor_adr': False,
'condim_max': False,
'dof_Madr': False,
'dof_armature': True,
'dof_bodyid': False,
@@ -1215,7 +1214,7 @@ _BATCH_DIM = {
'light_type': True,
'mapM2M': False,
'mat_rgba': True,
'mat_texid': True,
'mat_texid': False,
'mat_texrepeat': True,
'mesh_face': False,
'mesh_faceadr': False,
@@ -1249,6 +1248,7 @@ _BATCH_DIM = {
'nflexedge': False,
'nflexelem': False,
'nflexelemdata': False,
'nflexelemedge': False,
'nflexvert': False,
'ngeom': False,
'ngravcomp': False,
@@ -1257,16 +1257,22 @@ _BATCH_DIM = {
'njnt': False,
'nlight': False,
'nmat': False,
'nmaxcondim': False,
'nmaxmeshdeg': False,
'nmaxpolygon': False,
'nmaxpyramid': False,
'nmesh': False,
'nmeshface': False,
'nmeshgraph': False,
'nmeshnormal': False,
'nmeshpoly': False,
'nmeshpolymap': False,
'nmeshpolyvert': False,
'nmeshvert': False,
'nmocap': False,
'noct': False,
'npair': False,
'nplugin': False,
'nq': False,
'nrangefinder': False,
'nsensor': False,
@@ -1286,8 +1292,6 @@ _BATCH_DIM = {
'oct_aabb': False,
'oct_child': False,
'oct_coeff': False,
'opt__broadphase': False,
'opt__broadphase_filter': False,
'opt__ccd_iterations': False,
'opt__ccd_tolerance': True,
'opt__cone': False,
@@ -1405,8 +1409,6 @@ _BATCH_DIM = {
'wrap_type': False,
},
'Option': {
'broadphase': False,
'broadphase_filter': False,
'ccd_iterations': False,
'ccd_tolerance': True,
'cone': False,