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: 892971818
Change-Id: Ibce0c41294a4fcf2825a246a9533e8927fdc7e82
This commit is contained in:
Taylor Howell
2026-04-01 09:32:47 -07:00
committed by Copybara-Service
parent 70a7647ad9
commit e9de329e4e
36 changed files with 7231 additions and 4605 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
Vendored
+2
View File
@@ -64,6 +64,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.ray import rays as rays
from mujoco.mjx.third_party.mujoco_warp._src.render import render as render
from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_depth as get_depth
from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_rgb as get_rgb
from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_segmentation as get_segmentation
from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_pos as energy_pos
from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_vel as energy_vel
from mujoco.mjx.third_party.mujoco_warp._src.sensor import sensor_acc as sensor_acc
@@ -92,6 +93,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as x
from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType as BiasType
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter as BroadphaseFilter
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType as BroadphaseType
from mujoco.mjx.third_party.mujoco_warp._src.types import Callback as Callback
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType as ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import Constraint as Constraint
from mujoco.mjx.third_party.mujoco_warp._src.types import Contact as Contact
mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
Vendored
+366 -236
View File
@@ -189,12 +189,12 @@ def _compute_bvh_bounds(
upper_out: wp.array(dtype=wp.vec3),
group_out: wp.array(dtype=int),
):
world_id, geom_local_id = wp.tid()
worldid, geom_local_id = wp.tid()
geom_id = enabled_geom_ids[geom_local_id]
pos = geom_xpos_in[world_id, geom_id]
rot = geom_xmat_in[world_id, geom_id]
size = geom_size[world_id % geom_size.shape[0], geom_id]
pos = geom_xpos_in[worldid, geom_id]
rot = geom_xmat_in[worldid, geom_id]
size = geom_size[worldid % geom_size.shape[0], geom_id]
type = geom_type[geom_id]
# TODO: Investigate branch elimination with static loop unrolling
@@ -218,9 +218,9 @@ def _compute_bvh_bounds(
hfield_center = pos + rot[:, 2] * size[2]
lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size)
lower_out[world_id * bvh_ngeom + geom_local_id] = lower_bound
upper_out[world_id * bvh_ngeom + geom_local_id] = upper_bound
group_out[world_id * bvh_ngeom + geom_local_id] = world_id
lower_out[worldid * bvh_ngeom + geom_local_id] = lower_bound
upper_out[worldid * bvh_ngeom + geom_local_id] = upper_bound
group_out[worldid * bvh_ngeom + geom_local_id] = worldid
@wp.kernel
@@ -235,14 +235,70 @@ def compute_bvh_group_roots(
group_root_out[tid] = root
@wp.kernel
def _compute_flex_bvh_bounds(
# Model:
flex_vertadr: wp.array(dtype=int),
flex_vertnum: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_radius: wp.array(dtype=float),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
flex_geom_flexid: wp.array(dtype=int),
flex_geom_edgeid: wp.array(dtype=int),
bvh_ngeom: int,
total_bvh_size: int,
# Out:
lower_out: wp.array(dtype=wp.vec3),
upper_out: wp.array(dtype=wp.vec3),
group_out: wp.array(dtype=int),
):
worldid, flexlocalid = wp.tid()
flex_id = flex_geom_flexid[flexlocalid]
edge_id = flex_geom_edgeid[flexlocalid]
out_idx = worldid * total_bvh_size + bvh_ngeom + flexlocalid
radius = flex_radius[flex_id]
inflate = wp.vec3(radius, radius, radius)
if edge_id >= 0: # capsule (1D edge)
edge = flex_edge[edge_id]
vert_adr = flex_vertadr[flex_id]
v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
lower_out[out_idx] = wp.min(v0, v1) - inflate
upper_out[out_idx] = wp.max(v0, v1) + inflate
else: # mesh (2D/3D)
vert_adr = flex_vertadr[flex_id]
nvert = flex_vertnum[flex_id]
min_bound = wp.vec3(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL)
max_bound = wp.vec3(-MJ_MAXVAL, -MJ_MAXVAL, -MJ_MAXVAL)
for i in range(nvert):
v = flexvert_xpos_in[worldid, vert_adr + i]
min_bound = wp.min(min_bound, v)
max_bound = wp.max(max_bound, v)
lower_out[out_idx] = min_bound - inflate
upper_out[out_idx] = max_bound + inflate
group_out[out_idx] = worldid
def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext, nworld: int):
"""Build a global BVH for all geometries in all worlds."""
total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom
geom_type = wp.array(mjm.geom_type, dtype=int)
geom_dataid = wp.array(mjm.geom_dataid, dtype=int)
geom_size = wp.array(np.tile(mjm.geom_size[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
geom_xpos = wp.array(np.tile(mjd.geom_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
geom_xmat = wp.array(np.tile(mjd.geom_xmat.reshape(mjm.ngeom, 3, 3)[np.newaxis, :, :, :], (nworld, 1, 1, 1)), dtype=wp.mat33)
flex_vertadr = wp.array(mjm.flex_vertadr, dtype=int)
flex_vertnum = wp.array(mjm.flex_vertnum, dtype=int)
flex_edge = wp.array(mjm.flex_edge, dtype=wp.vec2i)
flex_radius = wp.array(mjm.flex_radius, dtype=float)
wp.launch(
kernel=_compute_bvh_bounds,
dim=(nworld, rc.bvh_ngeom),
@@ -252,7 +308,7 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
geom_size,
geom_xpos,
geom_xmat,
rc.bvh_ngeom,
total_bvh_size,
rc.enabled_geom_ids,
rc.mesh_bounds_size,
rc.hfield_bounds_size,
@@ -262,6 +318,26 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
],
)
flexvert_xpos = wp.array(np.tile(mjd.flexvert_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
wp.launch(
kernel=_compute_flex_bvh_bounds,
dim=(nworld, rc.bvh_nflexgeom),
inputs=[
flex_vertadr,
flex_vertnum,
flex_edge,
flex_radius,
flexvert_xpos,
rc.flex_geom_flexid,
rc.flex_geom_edgeid,
rc.bvh_ngeom,
total_bvh_size,
rc.lower,
rc.upper,
rc.group,
],
)
bvh = wp.Bvh(rc.lower, rc.upper, groups=rc.group, constructor="sah")
# BVH handle must be stored to avoid garbage collection
@@ -277,6 +353,8 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom
wp.launch(
kernel=_compute_bvh_bounds,
dim=(d.nworld, rc.bvh_ngeom),
@@ -286,7 +364,7 @@ def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
m.geom_size,
d.geom_xpos,
d.geom_xmat,
rc.bvh_ngeom,
total_bvh_size,
rc.enabled_geom_ids,
rc.mesh_bounds_size,
rc.hfield_bounds_size,
@@ -296,6 +374,26 @@ def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
],
)
if rc.bvh_nflexgeom > 0:
wp.launch(
kernel=_compute_flex_bvh_bounds,
dim=(d.nworld, rc.bvh_nflexgeom),
inputs=[
m.flex_vertadr,
m.flex_vertnum,
m.flex_edge,
m.flex_radius,
d.flexvert_xpos,
rc.flex_geom_flexid,
rc.flex_geom_edgeid,
rc.bvh_ngeom,
total_bvh_size,
rc.lower,
rc.upper,
rc.group,
],
)
rc.bvh.refit()
@@ -500,6 +598,12 @@ def build_hfield_bvh(
@wp.kernel
def accumulate_flex_vertex_normals(
# Model:
nflex: int,
flex_dim: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -509,10 +613,22 @@ def accumulate_flex_vertex_normals(
"""Accumulate per-vertex normals by summing adjacent face normals."""
worldid, elemid = wp.tid()
elem_base = elemid * 3
i0 = flex_elem[elem_base + 0]
i1 = flex_elem[elem_base + 1]
i2 = flex_elem[elem_base + 2]
for i in range(nflex):
locid = elemid - flex_elemadr[i]
if locid >= 0 and locid < flex_elemnum[i]:
f = i
break
if flex_dim[f] == 1 or flex_dim[f] == 3:
return
local_elemid = elemid - flex_elemadr[f]
elem_adr = flex_elemdataadr[f]
vert_adr = flex_vertadr[f]
elem_base = elem_adr + local_elemid * 3
i0 = vert_adr + flex_elem[elem_base + 0]
i1 = vert_adr + flex_elem[elem_base + 1]
i2 = vert_adr + flex_elem[elem_base + 2]
v0 = flexvert_xpos_in[worldid, i0]
v1 = flexvert_xpos_in[worldid, i1]
@@ -611,11 +727,12 @@ def _build_flex_2d_elements(
@wp.kernel
def _build_flex_2d_sides(
# Model:
flex_shell: wp.array(dtype=int),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
flexvert_norm_in: wp.array2d(dtype=wp.vec3),
flex_shell_in: wp.array(dtype=int),
shell_adr: int,
vert_adr: int,
face_offset: int,
@@ -635,8 +752,8 @@ def _build_flex_2d_sides(
worldid, shellid = wp.tid()
base = shell_adr + 2 * shellid
i0 = vert_adr + flex_shell_in[base + 0]
i1 = vert_adr + flex_shell_in[base + 1]
i0 = vert_adr + flex_shell[base + 0]
i1 = vert_adr + flex_shell[base + 1]
v0 = flexvert_xpos_in[worldid, i0]
v1 = flexvert_xpos_in[worldid, i1]
@@ -672,10 +789,11 @@ def _build_flex_2d_sides(
@wp.kernel
def _build_flex_3d_shells(
# Model:
flex_shell: wp.array(dtype=int),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
flex_shell_in: wp.array(dtype=int),
shell_adr: int,
vert_adr: int,
face_offset: int,
@@ -693,9 +811,9 @@ def _build_flex_3d_shells(
worldid, shellid = wp.tid()
base = shell_adr + shellid * 3
i0 = vert_adr + flex_shell_in[base + 0]
i1 = vert_adr + flex_shell_in[base + 1]
i2 = vert_adr + flex_shell_in[base + 2]
i0 = vert_adr + flex_shell[base + 0]
i1 = vert_adr + flex_shell[base + 1]
i2 = vert_adr + flex_shell[base + 2]
face_id = worldid * nface + face_offset + shellid
base = face_id * 3
@@ -716,163 +834,163 @@ def _build_flex_3d_shells(
@wp.kernel
def _update_flex_face_points(
def _update_flex_2d_face_points(
# Model:
nflex: int,
flex_dim: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_shell: wp.array(dtype=int),
flex_radius: wp.array(dtype=float),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
flex_shell_in: wp.array(dtype=int),
flexvert_norm_in: wp.array2d(dtype=wp.vec3),
flex_elemdataadr: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_faceadr: wp.array(dtype=int),
flex_radius: wp.array(dtype=float),
flex_workadr: wp.array(dtype=int),
flex_worknum: wp.array(dtype=int),
nfaces: int,
flex_id: int,
nface: int,
smooth: bool,
# Out:
face_point_out: wp.array(dtype=wp.vec3),
):
worldid, workid = wp.tid()
# identify which flex this work item belongs to
f = int(0)
locid = int(0)
for i in range(nflex):
locid = workid - flex_workadr[i]
if locid >= 0 and locid < flex_worknum[i]:
f = i
break
elem_adr = flex_elemdataadr[flex_id]
vert_adr = flex_vertadr[flex_id]
radius = flex_radius[flex_id]
nelem = flex_elemnum[flex_id]
world_face_offset = worldid * nface
dim = flex_dim[f]
face_offset = flex_faceadr[f]
world_face_offset = worldid * nfaces
vert_adr = flex_vertadr[f]
if dim == 2:
radius = flex_radius[f]
elem_count = flex_elemnum[f]
if locid < elem_count:
# 2D element faces
elemid = locid
elem_adr = flex_elemdataadr[f]
ebase = elem_adr + elemid * 3
i0 = vert_adr + flex_elem[ebase + 0]
i1 = vert_adr + flex_elem[ebase + 1]
i2 = vert_adr + flex_elem[ebase + 2]
v0 = flexvert_xpos_in[worldid, i0]
v1 = flexvert_xpos_in[worldid, i1]
v2 = flexvert_xpos_in[worldid, i2]
# TODO: Use static conditional
if smooth:
n0 = flexvert_norm_in[worldid, i0]
n1 = flexvert_norm_in[worldid, i1]
n2 = flexvert_norm_in[worldid, i2]
else:
face_nrm = wp.cross(v1 - v0, v2 - v0)
face_nrm = wp.normalize(face_nrm)
n0 = face_nrm
n1 = face_nrm
n2 = face_nrm
p0_pos = v0 + radius * n0
p1_pos = v1 + radius * n1
p2_pos = v2 + radius * n2
p0_neg = v0 - radius * n0
p1_neg = v1 - radius * n1
p2_neg = v2 - radius * n2
face_id0 = world_face_offset + face_offset + (2 * elemid)
base0 = face_id0 * 3
face_point_out[base0 + 0] = p0_pos
face_point_out[base0 + 1] = p1_pos
face_point_out[base0 + 2] = p2_pos
face_id1 = world_face_offset + face_offset + (2 * elemid + 1)
base1 = face_id1 * 3
face_point_out[base1 + 0] = p0_neg
face_point_out[base1 + 1] = p1_neg
face_point_out[base1 + 2] = p2_neg
else:
# 2D shell faces
shellid = locid - elem_count
shell_adr = flex_shelldataadr[f]
sbase = shell_adr + 2 * shellid
i0 = vert_adr + flex_shell_in[sbase + 0]
i1 = vert_adr + flex_shell_in[sbase + 1]
v0 = flexvert_xpos_in[worldid, i0]
v1 = flexvert_xpos_in[worldid, i1]
n0 = flexvert_norm_in[worldid, i0]
n1 = flexvert_norm_in[worldid, i1]
shell_face_offset = face_offset + (2 * elem_count)
face_id0 = world_face_offset + shell_face_offset + (2 * shellid)
base0 = face_id0 * 3
face_point_out[base0 + 0] = v0 + radius * n0
face_point_out[base0 + 1] = v1 - radius * n1
face_point_out[base0 + 2] = v1 + radius * n1
face_id1 = world_face_offset + shell_face_offset + (2 * shellid + 1)
base1 = face_id1 * 3
face_point_out[base1 + 0] = v1 - radius * n1
face_point_out[base1 + 1] = v0 + radius * n0
face_point_out[base1 + 2] = v0 - radius * n0
else:
# 3D shell faces
shellid = locid
shell_adr = flex_shelldataadr[f]
sbase = shell_adr + shellid * 3
i0 = vert_adr + flex_shell_in[sbase + 0]
i1 = vert_adr + flex_shell_in[sbase + 1]
i2 = vert_adr + flex_shell_in[sbase + 2]
if workid < nelem:
# 2D element faces
elemid = workid
ebase = elem_adr + elemid * 3
i0 = vert_adr + flex_elem[ebase + 0]
i1 = vert_adr + flex_elem[ebase + 1]
i2 = vert_adr + flex_elem[ebase + 2]
v0 = flexvert_xpos_in[worldid, i0]
v1 = flexvert_xpos_in[worldid, i1]
v2 = flexvert_xpos_in[worldid, i2]
face_id = world_face_offset + face_offset + shellid
fbase = face_id * 3
# TODO: Use static conditional
if smooth:
n0 = flexvert_norm_in[worldid, i0]
n1 = flexvert_norm_in[worldid, i1]
n2 = flexvert_norm_in[worldid, i2]
else:
face_nrm = wp.cross(v1 - v0, v2 - v0)
face_nrm = wp.normalize(face_nrm)
n0 = face_nrm
n1 = face_nrm
n2 = face_nrm
face_point_out[fbase + 0] = v0
face_point_out[fbase + 1] = v1
face_point_out[fbase + 2] = v2
p0_pos = v0 + radius * n0
p1_pos = v1 + radius * n1
p2_pos = v2 + radius * n2
p0_neg = v0 - radius * n0
p1_neg = v1 - radius * n1
p2_neg = v2 - radius * n2
face_id0 = world_face_offset + (2 * elemid)
base0 = face_id0 * 3
face_point_out[base0 + 0] = p0_pos
face_point_out[base0 + 1] = p1_pos
face_point_out[base0 + 2] = p2_pos
face_id1 = world_face_offset + (2 * elemid + 1)
base1 = face_id1 * 3
face_point_out[base1 + 0] = p0_neg
face_point_out[base1 + 1] = p1_neg
face_point_out[base1 + 2] = p2_neg
else:
# 2D shell faces
shell_adr = flex_shelldataadr[flex_id]
shellid = workid - nelem
sbase = shell_adr + 2 * shellid
i0 = vert_adr + flex_shell[sbase + 0]
i1 = vert_adr + flex_shell[sbase + 1]
v0 = flexvert_xpos_in[worldid, i0]
v1 = flexvert_xpos_in[worldid, i1]
n0 = flexvert_norm_in[worldid, i0]
n1 = flexvert_norm_in[worldid, i1]
shell_face_offset = 2 * nelem
face_id0 = world_face_offset + shell_face_offset + (2 * shellid)
base0 = face_id0 * 3
face_point_out[base0 + 0] = v0 + radius * n0
face_point_out[base0 + 1] = v1 - radius * n1
face_point_out[base0 + 2] = v1 + radius * n1
face_id1 = world_face_offset + shell_face_offset + (2 * shellid + 1)
base1 = face_id1 * 3
face_point_out[base1 + 0] = v1 - radius * n1
face_point_out[base1 + 1] = v0 + radius * n0
face_point_out[base1 + 2] = v0 - radius * n0
@wp.kernel
def _update_flex_3d_face_points(
# Model:
flex_vertadr: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_shell: wp.array(dtype=int),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
flex_id: int,
nface: int,
# Out:
face_point_out: wp.array(dtype=wp.vec3),
):
worldid, shellid = wp.tid()
shell_adr = flex_shelldataadr[flex_id]
vert_adr = flex_vertadr[flex_id]
face_id = worldid * nface + shellid
fbase = face_id * 3
sbase = shell_adr + shellid * 3
i0 = vert_adr + flex_shell[sbase + 0]
i1 = vert_adr + flex_shell[sbase + 1]
i2 = vert_adr + flex_shell[sbase + 2]
face_point_out[fbase + 0] = flexvert_xpos_in[worldid, i0]
face_point_out[fbase + 1] = flexvert_xpos_in[worldid, i1]
face_point_out[fbase + 2] = flexvert_xpos_in[worldid, i2]
def build_flex_bvh(
mjm: mujoco.MjModel, mjd: mujoco.MjData, nworld: int, constructor: str = "sah", leaf_size: int = 2
) -> tuple[wp.Mesh, wp.array, wp.array, wp.array, wp.array, wp.array, int]:
"""Create a Warp mesh BVH from flex data."""
if (mjm.flex_dim == 1).any():
raise ValueError("1D Flex objects are not currently supported.")
nflex = mjm.nflex
mjm: mujoco.MjModel,
mjd: mujoco.MjData,
nworld: int,
flex_id: int,
constructor: str = "sah",
leaf_size: int = 2,
) -> tuple[wp.Mesh, wp.array, wp.array, wp.array, int]:
"""Create a Warp mesh BVH for a single 2D or 3D flex."""
nflexvert = mjm.nflexvert
nflexelemdata = len(mjm.flex_elem)
flex_dim = wp.array(mjm.flex_dim, dtype=int)
flex_elemadr = wp.array(mjm.flex_elemadr, dtype=int)
flex_elemnum = wp.array(mjm.flex_elemnum, dtype=int)
flex_elem = wp.array(mjm.flex_elem, dtype=int)
flex_elemdataadr = wp.array(mjm.flex_elemdataadr, dtype=int)
flex_vertadr = wp.array(mjm.flex_vertadr, dtype=int)
flexvert_xpos = wp.array(np.tile(mjd.flexvert_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
flex_faceadr = [0]
for f in range(nflex):
if mjm.flex_dim[f] == 2:
flex_faceadr.append(flex_faceadr[-1] + 2 * mjm.flex_elemnum[f] + 2 * mjm.flex_shellnum[f])
elif mjm.flex_dim[f] == 3:
flex_faceadr.append(flex_faceadr[-1] + mjm.flex_shellnum[f])
dim = int(mjm.flex_dim[flex_id])
nelem = int(mjm.flex_elemnum[flex_id])
nshell = int(mjm.flex_shellnum[flex_id])
nface = int(flex_faceadr[-1])
flex_faceadr = flex_faceadr[:-1]
if dim == 2:
nface = 2 * nelem + 2 * nshell
else:
nface = nshell
face_point = wp.empty(nface * 3 * nworld, dtype=wp.vec3)
face_index = wp.empty(nface * 3 * nworld, dtype=wp.int32)
@@ -883,8 +1001,8 @@ def build_flex_bvh(
wp.launch(
kernel=accumulate_flex_vertex_normals,
dim=(nworld, nflexelemdata // 3),
inputs=[flex_elem, flexvert_xpos],
dim=(nworld, mjm.nflexelem),
inputs=[mjm.nflex, flex_dim, flex_vertadr, flex_elemadr, flex_elemnum, flex_elemdataadr, flex_elem, flexvert_xpos],
outputs=[flexvert_norm],
)
@@ -894,60 +1012,56 @@ def build_flex_bvh(
inputs=[flexvert_norm],
)
for f in range(nflex):
dim = mjm.flex_dim[f]
elem_adr = mjm.flex_elemdataadr[f]
nelem = mjm.flex_elemnum[f]
shell_adr = mjm.flex_shelldataadr[f]
nshell = mjm.flex_shellnum[f]
vert_adr = mjm.flex_vertadr[f]
elem_adr = mjm.flex_elemdataadr[flex_id]
shell_adr = mjm.flex_shelldataadr[flex_id]
vert_adr = mjm.flex_vertadr[flex_id]
if dim == 2:
wp.launch(
kernel=_build_flex_2d_elements,
dim=(nworld, nelem),
inputs=[
flex_elem,
flexvert_xpos,
flexvert_norm,
elem_adr,
vert_adr,
flex_faceadr[f],
mjm.flex_radius[f],
nface,
],
outputs=[face_point, face_index, group],
)
if dim == 2:
wp.launch(
kernel=_build_flex_2d_elements,
dim=(nworld, nelem),
inputs=[
flex_elem,
flexvert_xpos,
flexvert_norm,
elem_adr,
vert_adr,
0, # face_offset
mjm.flex_radius[flex_id],
nface,
],
outputs=[face_point, face_index, group],
)
wp.launch(
kernel=_build_flex_2d_sides,
dim=(nworld, nshell),
inputs=[
flexvert_xpos,
flexvert_norm,
flex_shell,
shell_adr,
vert_adr,
flex_faceadr[f] + 2 * nelem,
mjm.flex_radius[f],
nface,
],
outputs=[face_point, face_index, group],
)
elif dim == 3:
wp.launch(
kernel=_build_flex_3d_shells,
dim=(nworld, nshell),
inputs=[
flexvert_xpos,
flex_shell,
shell_adr,
vert_adr,
flex_faceadr[f],
nface,
],
outputs=[face_point, face_index, group],
)
wp.launch(
kernel=_build_flex_2d_sides,
dim=(nworld, nshell),
inputs=[
flex_shell,
flexvert_xpos,
flexvert_norm,
shell_adr,
vert_adr,
2 * nelem, # face_offset
mjm.flex_radius[flex_id],
nface,
],
outputs=[face_point, face_index, group],
)
elif dim == 3:
wp.launch(
kernel=_build_flex_3d_shells,
dim=(nworld, nshell),
inputs=[
flex_shell,
flexvert_xpos,
shell_adr,
vert_adr,
0, # face_offset
nface,
],
outputs=[face_point, face_index, group],
)
flex_mesh = wp.Mesh(
points=face_point,
@@ -965,24 +1079,23 @@ def build_flex_bvh(
outputs=[group_root],
)
return (
flex_mesh,
face_point,
group_root,
flex_shell,
flex_faceadr,
nface,
)
return flex_mesh, group_root
def refit_flex_bvh(m: Model, d: Data, rc: RenderContext):
"""Refit the flex BVH."""
"""Refit per-flex BVHs."""
flexvert_norm = wp.zeros(d.flexvert_xpos.shape, dtype=wp.vec3)
wp.launch(
kernel=accumulate_flex_vertex_normals,
dim=(d.nworld, m.nflexelemdata // 3),
dim=(d.nworld, m.nflexelem),
inputs=[
m.nflex,
m.flex_dim,
m.flex_vertadr,
m.flex_elemadr,
m.flex_elemnum,
m.flex_elemdataadr,
m.flex_elem,
d.flexvert_xpos,
],
@@ -991,32 +1104,49 @@ def refit_flex_bvh(m: Model, d: Data, rc: RenderContext):
wp.launch(
kernel=normalize_vertex_normals,
dim=(d.nworld, m.nflexvert),
dim=(d.nworld, d.flexvert_xpos.shape[1]),
inputs=[flexvert_norm],
)
wp.launch(
kernel=_update_flex_face_points,
dim=(d.nworld, rc.flex_nwork),
inputs=[
m.nflex,
m.flex_dim,
m.flex_vertadr,
m.flex_elemnum,
m.flex_elem,
d.flexvert_xpos,
rc.flex_shell,
flexvert_norm,
rc.flex_elemdataadr,
rc.flex_shelldataadr,
rc.flex_faceadr,
rc.flex_radius,
rc.flex_workadr,
rc.flex_worknum,
rc.flex_nface,
rc.flex_render_smooth,
],
outputs=[rc.flex_face_point],
)
for i in range(m.nflex):
if rc.flex_dim_np[i] == 1:
continue
mesh = rc.flex_mesh_registry[i]
nface = mesh.points.shape[0] // (3 * d.nworld)
rc.flex_mesh.refit()
if rc.flex_dim_np[i] == 2:
wp.launch(
kernel=_update_flex_2d_face_points,
dim=(d.nworld, nface // 2),
inputs=[
m.flex_vertadr,
m.flex_elemnum,
m.flex_elemdataadr,
m.flex_shelldataadr,
m.flex_elem,
m.flex_shell,
m.flex_radius,
d.flexvert_xpos,
flexvert_norm,
i,
nface,
rc.flex_render_smooth,
],
outputs=[mesh.points],
)
else:
wp.launch(
kernel=_update_flex_3d_face_points,
dim=(d.nworld, nface),
inputs=[
m.flex_vertadr,
m.flex_shelldataadr,
m.flex_shell,
d.flexvert_xpos,
i,
nface,
],
outputs=[mesh.points],
)
mesh.refit()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+19 -8
View File
@@ -15,34 +15,35 @@
from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
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 contact_params
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
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
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 mat43
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
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
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
import warp as wp
# TODO(team): improve compile time to enable backward pass
wp.set_module_options({"enable_backward": False})
@@ -233,6 +234,7 @@ def ccd_hfield_kernel_builder(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -516,6 +518,7 @@ def ccd_hfield_kernel_builder(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -572,6 +575,7 @@ def ccd_hfield_kernel_builder(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -626,6 +630,7 @@ def ccd_hfield_kernel_builder(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -681,6 +686,7 @@ def ccd_hfield_kernel_builder(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -751,6 +757,7 @@ def ccd_kernel_builder(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -871,6 +878,7 @@ def ccd_kernel_builder(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -956,6 +964,7 @@ def ccd_kernel_builder(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1070,6 +1079,7 @@ def ccd_kernel_builder(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -1156,6 +1166,7 @@ def convex_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+102 -111
View File
@@ -18,14 +18,15 @@
import dataclasses
from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
import warp as wp
wp.set_module_options({"enable_backward": False})
@@ -63,30 +64,30 @@ class Geom:
@wp.func
def geom_collision_pair(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graph: wp.array(dtype=int),
mesh_polynum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
mesh_polynormal: wp.array(dtype=wp.vec3),
mesh_polyvertadr: wp.array(dtype=int),
mesh_polyvertnum: wp.array(dtype=int),
mesh_polyvert: wp.array(dtype=int),
mesh_polymapadr: wp.array(dtype=int),
mesh_polymapnum: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
geoms: wp.vec2i,
worldid: int,
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graph: wp.array(dtype=int),
mesh_polynum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
mesh_polynormal: wp.array(dtype=wp.vec3),
mesh_polyvertadr: wp.array(dtype=int),
mesh_polyvertnum: wp.array(dtype=int),
mesh_polyvert: wp.array(dtype=int),
mesh_polymapadr: wp.array(dtype=int),
mesh_polymapnum: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
geoms: wp.vec2i,
worldid: int,
) -> Tuple[Geom, Geom]:
geom1 = Geom()
geom2 = Geom()
@@ -155,38 +156,39 @@ def geom_collision_pair(
@wp.func
def write_contact(
# Data in:
naconmax_in: int,
# In:
id_: int,
dist_in: float,
pos_in: wp.vec3,
frame_in: wp.mat33,
margin_in: float,
gap_in: float,
condim_in: int,
friction_in: vec5,
solref_in: wp.vec2,
solreffriction_in: wp.vec2,
solimp_in: vec5,
geoms_in: wp.vec2i,
pairid_in: wp.vec2i,
worldid_in: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
# Data in:
naconmax_in: int,
# In:
id_: int,
dist_in: float,
pos_in: wp.vec3,
frame_in: wp.mat33,
margin_in: float,
gap_in: float,
condim_in: int,
friction_in: vec5,
solref_in: wp.vec2,
solreffriction_in: wp.vec2,
solimp_in: vec5,
geoms_in: wp.vec2i,
pairid_in: wp.vec2i,
worldid_in: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
) -> int:
"""Atomically write a detected contact into the contact output arrays.
@@ -222,33 +224,35 @@ def write_contact(
contact_solimp_out[cid] = solimp_in
contact_type_out[cid] = contact_type
contact_geomcollisionid_out[cid] = id_
for i in range(contact_efc_address_out.shape[1]):
contact_efc_address_out[cid, i] = -1
return int(active)
return 0
@wp.func
def contact_params(
# Model:
geom_condim: wp.array(dtype=int),
geom_priority: wp.array(dtype=int),
geom_solmix: wp.array2d(dtype=float),
geom_solref: wp.array2d(dtype=wp.vec2),
geom_solimp: wp.array2d(dtype=vec5),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
geom_gap: wp.array2d(dtype=float),
pair_dim: wp.array(dtype=int),
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
pair_solimp: wp.array2d(dtype=vec5),
pair_margin: wp.array2d(dtype=float),
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
cid: int,
worldid: int,
# Model:
geom_condim: wp.array(dtype=int),
geom_priority: wp.array(dtype=int),
geom_solmix: wp.array2d(dtype=float),
geom_solref: wp.array2d(dtype=wp.vec2),
geom_solimp: wp.array2d(dtype=vec5),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
geom_gap: wp.array2d(dtype=float),
pair_dim: wp.array(dtype=int),
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
pair_solimp: wp.array2d(dtype=vec5),
pair_margin: wp.array2d(dtype=float),
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
cid: int,
worldid: int,
):
"""Resolve contact parameters for a collision pair.
@@ -267,9 +271,7 @@ def contact_params(
condim = pair_dim[pairid]
friction = pair_friction[worldid % pair_friction.shape[0], pairid]
solref = pair_solref[worldid % pair_solref.shape[0], pairid]
solreffriction = pair_solreffriction[
worldid % pair_solreffriction.shape[0], pairid
]
solreffriction = pair_solreffriction[worldid % pair_solreffriction.shape[0], pairid]
solimp = pair_solimp[worldid % pair_solimp.shape[0], pairid]
else:
g1 = geoms[0]
@@ -305,44 +307,33 @@ def contact_params(
mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
condim = wp.max(condim1, condim2)
max_geom_friction = wp.max(
geom_friction[friction_id, g1], geom_friction[friction_id, g2]
)
max_geom_friction = wp.max(geom_friction[friction_id, g1], geom_friction[friction_id, g2])
friction = vec5(
max_geom_friction[0],
max_geom_friction[0],
max_geom_friction[1],
max_geom_friction[2],
max_geom_friction[2],
max_geom_friction[0],
max_geom_friction[0],
max_geom_friction[1],
max_geom_friction[2],
max_geom_friction[2],
)
if (
geom_solref[solref_id, g1][0] > 0.0
and geom_solref[solref_id, g2][0] > 0.0
):
solref = (
mix * geom_solref[solref_id, g1]
+ (1.0 - mix) * geom_solref[solref_id, g2]
)
if geom_solref[solref_id, g1][0] > 0.0 and geom_solref[solref_id, g2][0] > 0.0:
solref = mix * geom_solref[solref_id, g1] + (1.0 - mix) * geom_solref[solref_id, g2]
else:
solref = wp.min(geom_solref[solref_id, g1], geom_solref[solref_id, g2])
solreffriction = wp.vec2(0.0, 0.0)
solimp = (
mix * geom_solimp[solimp_id, g1]
+ (1.0 - mix) * geom_solimp[solimp_id, g2]
)
solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
# geom priority is ignored
margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
friction = vec5(
wp.max(MJ_MINMU, friction[0]),
wp.max(MJ_MINMU, friction[1]),
wp.max(MJ_MINMU, friction[2]),
wp.max(MJ_MINMU, friction[3]),
wp.max(MJ_MINMU, friction[4]),
wp.max(MJ_MINMU, friction[0]),
wp.max(MJ_MINMU, friction[1]),
wp.max(MJ_MINMU, friction[2]),
wp.max(MJ_MINMU, friction[3]),
wp.max(MJ_MINMU, friction[4]),
)
return geoms, margin, gap, condim, friction, solref, solreffriction, solimp
@@ -366,7 +357,7 @@ class CollisionContext:
def create_collision_context(naconmax: int) -> CollisionContext:
"""Create a CollisionContext with allocated arrays."""
return CollisionContext(
collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
collision_pairid=wp.empty(naconmax, dtype=wp.vec2i),
collision_worldid=wp.empty(naconmax, dtype=int),
collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
collision_pairid=wp.empty(naconmax, dtype=wp.vec2i),
collision_worldid=wp.empty(naconmax, dtype=int),
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+12 -19
View File
@@ -15,25 +15,27 @@
from typing import Any
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_convex import convex_narrowphase
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import create_collision_context
from mujoco.mjx.third_party.mujoco_warp._src.collision_flex import flex_narrowphase
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import primitive_narrowphase
from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf_narrowphase
from mujoco.mjx.third_party.mujoco_warp._src.math import upper_tri_index
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType
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 GeomType
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.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
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
import warp as wp
wp.set_module_options({"enable_backward": False})
@@ -290,25 +292,13 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
# 8: obb
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],
) # kernel_analyzer: ignore
size1, size2 = (
geom_aabb[aabb_id, geom1, 1],
geom_aabb[aabb_id, geom2, 1],
) # kernel_analyzer: ignore
center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0] # kernel_analyzer: ignore
size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1] # kernel_analyzer: ignore
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],
) # kernel_analyzer: ignore
rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2] # kernel_analyzer: ignore
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],
) # kernel_analyzer: ignore
margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] # kernel_analyzer: ignore
xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
@@ -757,6 +747,9 @@ def _narrowphase(m: Model, d: Data, ctx: CollisionContext):
if m.has_sdf_geom:
sdf_narrowphase(m, d, ctx)
if m.nflex > 0:
flex_narrowphase(m, d)
@event_scope
def collision(m: Model, d: Data):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py
Vendored
+834
View File
@@ -0,0 +1,834 @@
# Copyright 2026 The Newton Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Flex collision detection (geom vs flex triangles)."""
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import collision_primitive_core
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
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.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
@wp.func
def _write_flex_contact(
# Data in:
naconmax_in: int,
# In:
dist: float,
pos: wp.vec3,
frame: wp.mat33,
margin: float,
condim: int,
friction: vec5,
solref: wp.vec2,
solimp: vec5,
geom: int,
flexid: int,
vertid: int,
worldid: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_flex_out: wp.array(dtype=wp.vec2i),
contact_vert_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
if dist >= margin or dist >= MJ_MAXVAL:
return
id_ = wp.atomic_add(nacon_out, 0, 1)
if id_ >= naconmax_in:
return
contact_dist_out[id_] = dist
contact_pos_out[id_] = pos
contact_frame_out[id_] = frame
contact_includemargin_out[id_] = margin
contact_friction_out[id_] = friction
contact_solref_out[id_] = solref
contact_solreffriction_out[id_] = wp.vec2(0.0, 0.0)
contact_solimp_out[id_] = solimp
contact_dim_out[id_] = condim
contact_geom_out[id_] = wp.vec2i(geom, -1)
contact_flex_out[id_] = wp.vec2i(-1, flexid)
contact_vert_out[id_] = wp.vec2i(-1, vertid)
contact_worldid_out[id_] = worldid
contact_type_out[id_] = 1
contact_geomcollisionid_out[id_] = 0
@wp.func
def _collide_geom_triangle(
# Data in:
naconmax_in: int,
# In:
gtype: int,
pos: wp.vec3,
rot: wp.mat33,
size_val: wp.vec3,
t1: wp.vec3,
t2: wp.vec3,
t3: wp.vec3,
tri_radius: float,
margin: float,
condim: int,
friction: vec5,
solref: wp.vec2,
solimp: vec5,
geomid: int,
flexid: int,
vertex_id: int,
worldid: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_flex_out: wp.array(dtype=wp.vec2i),
contact_vert_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
if gtype == int(GeomType.SPHERE):
sphere_radius = size_val[0]
dist, contact_pos, nrm = collision_primitive_core.sphere_triangle(pos, sphere_radius, t1, t2, t3, tri_radius)
if dist < margin:
_write_flex_contact(
naconmax_in,
dist,
contact_pos,
make_frame(nrm),
margin,
condim,
friction,
solref,
solimp,
geomid,
flexid,
vertex_id,
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_flex_out,
contact_vert_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
return
# Capsule, box, cylinder all return up to 2 contacts - compute then share writing code
dists = wp.vec2(collision_primitive_core.MJ_MAXVAL, collision_primitive_core.MJ_MAXVAL)
poss = collision_primitive_core.mat23f(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
nrms = collision_primitive_core.mat23f(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
if gtype == int(GeomType.CAPSULE):
cap_radius = size_val[0]
cap_half_len = size_val[1]
cap_axis = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2])
dists, poss, nrms = collision_primitive_core.capsule_triangle(
pos, cap_axis, cap_radius, cap_half_len, t1, t2, t3, tri_radius
)
elif gtype == int(GeomType.BOX):
dists, poss, nrms = collision_primitive_core.box_triangle(pos, rot, size_val, t1, t2, t3, tri_radius)
elif gtype == int(GeomType.CYLINDER):
cyl_radius = size_val[0]
cyl_half_height = size_val[1]
cyl_axis = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2])
dists, poss, nrms = collision_primitive_core.cylinder_triangle(
pos, cyl_axis, cyl_radius, cyl_half_height, t1, t2, t3, tri_radius
)
# Write up to 2 contacts (shared code for capsule/box/cylinder)
if dists[0] < margin:
p1 = wp.vec3(poss[0, 0], poss[0, 1], poss[0, 2])
n1 = wp.vec3(nrms[0, 0], nrms[0, 1], nrms[0, 2])
_write_flex_contact(
naconmax_in,
dists[0],
p1,
make_frame(n1),
margin,
condim,
friction,
solref,
solimp,
geomid,
flexid,
vertex_id,
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_flex_out,
contact_vert_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
if dists[1] < margin:
p2 = wp.vec3(poss[1, 0], poss[1, 1], poss[1, 2])
n2 = wp.vec3(nrms[1, 0], nrms[1, 1], nrms[1, 2])
_write_flex_contact(
naconmax_in,
dists[1],
p2,
make_frame(n2),
margin,
condim,
friction,
solref,
solimp,
geomid,
flexid,
vertex_id,
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_flex_out,
contact_vert_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
@wp.kernel
def _flex_plane_narrowphase(
# Model:
ngeom: int,
nflexvert: int,
geom_type: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_solref: wp.array2d(dtype=wp.vec2),
geom_solimp: wp.array2d(dtype=vec5),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
flex_condim: wp.array(dtype=int),
flex_friction: wp.array(dtype=wp.vec3),
flex_margin: wp.array(dtype=float),
flex_vertadr: wp.array(dtype=int),
flex_radius: wp.array(dtype=float),
flex_vertflexid: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
nworld_in: int,
naconmax_in: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_flex_out: wp.array(dtype=wp.vec2i),
contact_vert_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
worldid, vertid = wp.tid()
flexid = flex_vertflexid[vertid]
radius = flex_radius[flexid]
flex_margin_val = flex_margin[flexid]
flex_condim_val = flex_condim[flexid]
flex_fric = flex_friction[flexid]
# Convert global vertid to local vertex index within this flex
local_vertid = vertid - flex_vertadr[flexid]
vert = flexvert_xpos_in[worldid, vertid]
# TODO: Add a broadphase
for geomid in range(ngeom):
gtype = geom_type[geomid]
if gtype != int(GeomType.PLANE):
continue
plane_pos = geom_xpos_in[worldid, geomid]
plane_rot = geom_xmat_in[worldid, geomid]
plane_normal = wp.vec3(plane_rot[0, 2], plane_rot[1, 2], plane_rot[2, 2])
margin = geom_margin[worldid % geom_margin.shape[0], geomid] + flex_margin_val
diff = vert - plane_pos
signed_dist = wp.dot(diff, plane_normal)
dist = signed_dist - radius
if dist < margin:
geom_condim_val = geom_condim[geomid]
condim = wp.max(geom_condim_val, flex_condim_val)
solref = geom_solref[worldid % geom_solref.shape[0], geomid]
solimp = geom_solimp[worldid % geom_solimp.shape[0], geomid]
geom_fric = geom_friction[worldid % geom_friction.shape[0], geomid]
fric0 = wp.max(geom_fric[0], flex_fric[0])
fric1 = wp.max(geom_fric[1], flex_fric[1])
fric2 = wp.max(geom_fric[2], flex_fric[2])
friction = vec5(
wp.max(MJ_MINMU, fric0),
wp.max(MJ_MINMU, fric0),
wp.max(MJ_MINMU, fric1),
wp.max(MJ_MINMU, fric2),
wp.max(MJ_MINMU, fric2),
)
contact_pos = vert - plane_normal * (dist * 0.5 + radius)
_write_flex_contact(
naconmax_in,
dist,
contact_pos,
make_frame(plane_normal),
margin,
condim,
friction,
solref,
solimp,
geomid,
flexid,
local_vertid,
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_flex_out,
contact_vert_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
@wp.kernel
def _flex_narrowphase_dim2(
# Model:
ngeom: int,
nflex: int,
geom_type: wp.array(dtype=int),
geom_contype: wp.array(dtype=int),
geom_conaffinity: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_solref: wp.array2d(dtype=wp.vec2),
geom_solimp: wp.array2d(dtype=vec5),
geom_size: wp.array2d(dtype=wp.vec3),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
flex_contype: wp.array(dtype=int),
flex_conaffinity: wp.array(dtype=int),
flex_margin: wp.array(dtype=float),
flex_dim: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_radius: wp.array(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
nworld_in: int,
naconmax_in: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_flex_out: wp.array(dtype=wp.vec2i),
contact_vert_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
worldid, elemid = wp.tid()
flexid = int(-1)
for i in range(nflex):
if flex_dim[i] != 2:
continue
elem_adr = flex_elemadr[i]
elem_num = flex_elemnum[i]
if elemid >= elem_adr and elemid < elem_adr + elem_num:
flexid = i
break
if flexid < 0:
return
vert_adr = flex_vertadr[flexid]
tri_radius = flex_radius[flexid]
tri_margin = flex_margin[flexid]
elem_data_idx = flex_elemdataadr[flexid] + (elemid - flex_elemadr[flexid]) * 3
v0_local = flex_elem[elem_data_idx]
v1_local = flex_elem[elem_data_idx + 1]
v2_local = flex_elem[elem_data_idx + 2]
t1 = flexvert_xpos_in[worldid, vert_adr + v0_local]
t2 = flexvert_xpos_in[worldid, vert_adr + v1_local]
t3 = flexvert_xpos_in[worldid, vert_adr + v2_local]
# TODO: Add a broadphase
for geomid in range(ngeom):
gtype = geom_type[geomid]
if (
gtype != int(GeomType.SPHERE)
and gtype != int(GeomType.CAPSULE)
and gtype != int(GeomType.BOX)
and gtype != int(GeomType.CYLINDER)
):
continue
g_contype = geom_contype[geomid]
g_conaffinity = geom_conaffinity[geomid]
f_contype = flex_contype[flexid]
f_conaffinity = flex_conaffinity[flexid]
if not ((g_contype & f_conaffinity) or (f_contype & g_conaffinity)):
continue
geom_margin_val = geom_margin[worldid % geom_margin.shape[0], geomid]
margin = geom_margin_val + tri_margin
geom_pos = geom_xpos_in[worldid, geomid]
geom_rot = geom_xmat_in[worldid, geomid]
geom_size_val = geom_size[worldid % geom_size.shape[0], geomid]
condim = geom_condim[geomid]
gf = geom_friction[worldid % geom_friction.shape[0], geomid]
friction = vec5(
wp.max(MJ_MINMU, gf[0]),
wp.max(MJ_MINMU, gf[0]),
wp.max(MJ_MINMU, gf[1]),
wp.max(MJ_MINMU, gf[2]),
wp.max(MJ_MINMU, gf[2]),
)
solref = geom_solref[worldid % geom_solref.shape[0], geomid]
solimp = geom_solimp[worldid % geom_solimp.shape[0], geomid]
_collide_geom_triangle(
naconmax_in,
gtype,
geom_pos,
geom_rot,
geom_size_val,
t1,
t2,
t3,
tri_radius,
margin,
condim,
friction,
solref,
solimp,
geomid,
flexid,
v0_local,
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_flex_out,
contact_vert_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
@wp.kernel
def _flex_narrowphase_dim3(
# Model:
ngeom: int,
nflex: int,
geom_type: wp.array(dtype=int),
geom_contype: wp.array(dtype=int),
geom_conaffinity: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_solref: wp.array2d(dtype=wp.vec2),
geom_solimp: wp.array2d(dtype=vec5),
geom_size: wp.array2d(dtype=wp.vec3),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
flex_contype: wp.array(dtype=int),
flex_conaffinity: wp.array(dtype=int),
flex_margin: wp.array(dtype=float),
flex_dim: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_shellnum: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_shell: wp.array(dtype=int),
flex_radius: wp.array(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
nworld_in: int,
naconmax_in: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_flex_out: wp.array(dtype=wp.vec2i),
contact_vert_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
worldid, shellid = wp.tid()
flexid = int(-1)
shell_offset = int(0)
for i in range(nflex):
if flex_dim[i] != 3:
continue
shell_num = flex_shellnum[i]
if shellid >= shell_offset and shellid < shell_offset + shell_num:
flexid = i
break
shell_offset += shell_num
if flexid < 0:
return
vert_adr = flex_vertadr[flexid]
tri_radius = flex_radius[flexid]
tri_margin = flex_margin[flexid]
shell_adr = flex_shelldataadr[flexid]
local_shellid = shellid - shell_offset
shell_data_idx = shell_adr + local_shellid * 3
v0_local = flex_shell[shell_data_idx]
v1_local = flex_shell[shell_data_idx + 1]
v2_local = flex_shell[shell_data_idx + 2]
t1 = flexvert_xpos_in[worldid, vert_adr + v0_local]
t2 = flexvert_xpos_in[worldid, vert_adr + v1_local]
t3 = flexvert_xpos_in[worldid, vert_adr + v2_local]
# TODO: Add a broadphase
for geomid in range(ngeom):
gtype = geom_type[geomid]
if (
gtype != int(GeomType.SPHERE)
and gtype != int(GeomType.CAPSULE)
and gtype != int(GeomType.BOX)
and gtype != int(GeomType.CYLINDER)
):
continue
g_contype = geom_contype[geomid]
g_conaffinity = geom_conaffinity[geomid]
f_contype = flex_contype[flexid]
f_conaffinity = flex_conaffinity[flexid]
if not ((g_contype & f_conaffinity) or (f_contype & g_conaffinity)):
continue
geom_margin_val = geom_margin[worldid % geom_margin.shape[0], geomid]
margin = geom_margin_val + tri_margin
geom_pos = geom_xpos_in[worldid, geomid]
geom_rot = geom_xmat_in[worldid, geomid]
geom_size_val = geom_size[worldid % geom_size.shape[0], geomid]
condim = geom_condim[geomid]
gf = geom_friction[worldid % geom_friction.shape[0], geomid]
friction = vec5(
wp.max(MJ_MINMU, gf[0]),
wp.max(MJ_MINMU, gf[0]),
wp.max(MJ_MINMU, gf[1]),
wp.max(MJ_MINMU, gf[2]),
wp.max(MJ_MINMU, gf[2]),
)
solref = geom_solref[worldid % geom_solref.shape[0], geomid]
solimp = geom_solimp[worldid % geom_solimp.shape[0], geomid]
_collide_geom_triangle(
naconmax_in,
gtype,
geom_pos,
geom_rot,
geom_size_val,
t1,
t2,
t3,
tri_radius,
margin,
condim,
friction,
solref,
solimp,
geomid,
flexid,
v0_local,
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_flex_out,
contact_vert_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
nacon_out,
)
@event_scope
def flex_narrowphase(m: Model, d: Data):
"""Runs collision detection between geoms and flex elements."""
if m.nflex == 0:
return
wp.launch(
_flex_narrowphase_dim2,
dim=(d.nworld, m.nflexelem),
inputs=[
m.ngeom,
m.nflex,
m.geom_type,
m.geom_contype,
m.geom_conaffinity,
m.geom_condim,
m.geom_solref,
m.geom_solimp,
m.geom_size,
m.geom_friction,
m.geom_margin,
m.flex_contype,
m.flex_conaffinity,
m.flex_margin,
m.flex_dim,
m.flex_vertadr,
m.flex_elemadr,
m.flex_elemnum,
m.flex_elemdataadr,
m.flex_elem,
m.flex_radius,
d.geom_xpos,
d.geom_xmat,
d.flexvert_xpos,
d.nworld,
d.naconmax,
],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.flex,
d.contact.vert,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
],
)
wp.launch(
_flex_narrowphase_dim3,
dim=(d.nworld, m.nflexshelldata // 3),
inputs=[
m.ngeom,
m.nflex,
m.geom_type,
m.geom_contype,
m.geom_conaffinity,
m.geom_condim,
m.geom_solref,
m.geom_solimp,
m.geom_size,
m.geom_friction,
m.geom_margin,
m.flex_contype,
m.flex_conaffinity,
m.flex_margin,
m.flex_dim,
m.flex_vertadr,
m.flex_shellnum,
m.flex_shelldataadr,
m.flex_shell,
m.flex_radius,
d.geom_xpos,
d.geom_xmat,
d.flexvert_xpos,
d.nworld,
d.naconmax,
],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.flex,
d.contact.vert,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
],
)
wp.launch(
_flex_plane_narrowphase,
dim=(d.nworld, m.nflexvert),
inputs=[
m.ngeom,
m.nflexvert,
m.geom_type,
m.geom_condim,
m.geom_solref,
m.geom_solimp,
m.geom_friction,
m.geom_margin,
m.flex_condim,
m.flex_friction,
m.flex_margin,
m.flex_vertadr,
m.flex_radius,
m.flex_vertflexid,
d.geom_xpos,
d.geom_xmat,
d.flexvert_xpos,
d.nworld,
d.naconmax,
],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.flex,
d.contact.vert,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+10 -21
View File
@@ -16,11 +16,12 @@
import math
from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
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
import warp as wp
# TODO(team): improve compile time to enable backward pass
wp.set_module_options({"enable_backward": False})
@@ -581,16 +582,19 @@ def gjk(
simplex_index2 = wp.vec4i()
n = int(0)
coordinates = wp.vec4() # barycentric coordinates
epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance)
tol2 = tolerance * tolerance
epsilon = wp.where(is_discrete, 0.0, 0.5 * tol2)
# set initial guess
x_k = x1_0 - x2_0
xnorm_old = FLOAT_MAX
for k in range(gjk_iterations):
for _ in range(gjk_iterations):
xnorm = wp.dot(x_k, x_k)
# TODO(kbayes): determine new constant here
if xnorm < 1e-12:
if xnorm < tol2 or wp.abs(xnorm_old - xnorm) < tol2:
break
xnorm_old = xnorm
dir_neg = x_k / wp.sqrt(xnorm)
# compute kth support point in geom1
@@ -663,13 +667,6 @@ def gjk(
if n == 4:
break
if k == gjk_iterations - 1:
wp.printf(
"Warning: opt.ccd_iterations, currently set to %d, needs to be"
" increased.\n",
gjk_iterations,
)
result = GJKResult()
# compute the approximate witness points
@@ -1205,7 +1202,6 @@ def _is_invalid_face(face: int) -> bool:
def _epa(
# In:
tolerance: float,
gjk_iterations: int,
epa_iterations: int,
pt: Polytope,
geom1: Geom,
@@ -1226,7 +1222,7 @@ def _epa(
# 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 k in range(epa_iterations):
for _ in range(epa_iterations):
pidx = idx
idx = int(-1)
lower2 = float(FLOAT_MAX)
@@ -1325,13 +1321,6 @@ def _epa(
# clear horizon
pt.nhorizon = 0
if k == epa_iterations - 1:
wp.printf(
"Warning: opt.ccd_iterations, currently set to %d, needs to be"
" increased.\n",
gjk_iterations,
)
# return from valid face
if idx > -1:
x1, x2, dist = _epa_witness(pt, geom1, geom2, geomtype1, geomtype2, idx)
@@ -2347,7 +2336,7 @@ def ccd(
if pt.status:
return result.dist, 1, result.x1, result.x2, -1
dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
if idx == -1:
return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+34 -4
View File
@@ -15,9 +15,11 @@
from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import box_box
@@ -34,15 +36,14 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_sphere
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
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 mat43
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
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 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
import warp as wp
wp.set_module_options({"enable_backward": False})
@@ -304,6 +305,7 @@ def plane_sphere_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -339,6 +341,7 @@ def plane_sphere_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -374,6 +377,7 @@ def sphere_sphere_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -408,6 +412,7 @@ def sphere_sphere_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -443,6 +448,7 @@ def sphere_capsule_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -480,6 +486,7 @@ def sphere_capsule_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -515,6 +522,7 @@ def capsule_capsule_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -564,6 +572,7 @@ def capsule_capsule_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -599,6 +608,7 @@ def plane_capsule_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -644,6 +654,7 @@ def plane_capsule_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -679,6 +690,7 @@ def plane_ellipsoid_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -713,6 +725,7 @@ def plane_ellipsoid_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -748,6 +761,7 @@ def plane_box_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -784,6 +798,7 @@ def plane_box_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -819,6 +834,7 @@ def plane_convex_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -855,6 +871,7 @@ def plane_convex_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -890,6 +907,7 @@ def sphere_cylinder_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -934,6 +952,7 @@ def sphere_cylinder_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -969,6 +988,7 @@ def plane_cylinder_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1015,6 +1035,7 @@ def plane_cylinder_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -1050,6 +1071,7 @@ def sphere_box_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1083,6 +1105,7 @@ def sphere_box_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -1118,6 +1141,7 @@ def capsule_box_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1166,6 +1190,7 @@ def capsule_box_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -1201,6 +1226,7 @@ def box_box_wrapper(
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1245,6 +1271,7 @@ def box_box_wrapper(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -1327,6 +1354,7 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1416,6 +1444,7 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -1511,6 +1540,7 @@ def primitive_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_ta
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py
Vendored
+504
View File
@@ -1489,3 +1489,507 @@ def capsule_box(
mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
mat23f(normal1[0], normal1[1], normal1[2], normal2[0], normal2[1], normal2[2]),
)
@wp.func
def _tri_area_sign(p1: wp.vec2, p2: wp.vec2, p3: wp.vec2) -> float:
"""Sign of (signed) area of planar triangle."""
return wp.sign((p1[0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p3[1]))
@wp.func
def _tri_point_segment(p: wp.vec2, u: wp.vec2, v: wp.vec2) -> wp.vec2:
"""Find nearest point to p within line segment (u, v)."""
uv = v - u
up = p - u
denom = wp.max(MJ_MINVAL, wp.dot(uv, uv))
a = wp.dot(uv, up) / denom
if a <= 0.0:
return u
elif a >= 1.0:
return v
else:
return u + a * uv
@wp.func
def sphere_triangle(
sphere_pos: wp.vec3,
sphere_radius: float,
t1: wp.vec3,
t2: wp.vec3,
t3: wp.vec3,
tri_radius: float,
) -> Tuple[float, wp.vec3, wp.vec3]:
"""Core contact geometry calculation for sphere-triangle collision.
Port of mjraw_SphereTriangle from engine_collision_primitive.c
Args:
sphere_pos: Center position of the sphere.
sphere_radius: Radius of the sphere.
t1: Triangle vertex positions.
t2: Triangle vertex positions.
t3: Triangle vertex positions.
tri_radius: Triangle (flex element) radius.
Returns:
- Contact distance (MJ_MAXVAL if no collision).
- Contact position.
- Contact normal vector.
"""
S = sphere_pos - t1
A = t2 - t1
B = t3 - t1
N = wp.normalize(wp.cross(A, B))
dstS = wp.dot(N, S)
P = S - dstS * N
V1 = wp.normalize(A)
lenA = wp.length(A)
V2 = wp.normalize(wp.cross(N, A))
o = wp.vec2(0.0, 0.0)
a = wp.vec2(lenA, 0.0)
b = wp.vec2(wp.dot(V1, B), wp.dot(V2, B))
p = wp.vec2(wp.dot(V1, P), wp.dot(V2, P))
sign1 = _tri_area_sign(p, o, a)
sign2 = _tri_area_sign(p, a, b)
sign3 = _tri_area_sign(p, b, o)
X = wp.vec3(0.0)
if sign1 == sign2 and sign2 == sign3:
X = P
else:
x0 = _tri_point_segment(p, o, a)
x1 = _tri_point_segment(p, a, b)
x2 = _tri_point_segment(p, b, o)
d0 = wp.length(p - x0)
d1 = wp.length(p - x1)
d2 = wp.length(p - x2)
if d0 < d1 and d0 < d2:
X = x0[0] * V1 + x0[1] * V2
elif d1 < d2:
X = x1[0] * V1 + x1[1] * V2
else:
X = x2[0] * V1 + x2[1] * V2
nrm = X - S
dst = wp.length(nrm)
if dst > MJ_MINVAL:
nrm = nrm / dst
else:
nrm = N
dist = dst - sphere_radius - tri_radius
pos = sphere_pos + nrm * (sphere_radius + 0.5 * dist)
return dist, pos, nrm
@wp.func
def box_triangle(
box_pos: wp.vec3,
box_rot: wp.mat33,
box_size: wp.vec3,
t1: wp.vec3,
t2: wp.vec3,
t3: wp.vec3,
tri_radius: float,
) -> Tuple[wp.vec2, mat23f, mat23f]:
"""Core contact geometry calculation for box-triangle collision.
Port of mjraw_BoxTriangle from engine_collision_primitive.c
Args:
box_pos: Center position of the box.
box_rot: Orientation matrix of the box.
box_size: Half-sizes of the box.
t1: Triangle vertex positions.
t2: Triangle vertex positions.
t3: Triangle vertex positions.
tri_radius: Triangle (flex element) radius.
Returns:
- wp.vec2 of distances for up to 2 contacts (MJ_MAXVAL if no collision).
- mat23f of contact positions (2 x vec3).
- mat23f of contact normals (2 x vec3).
"""
dist1 = MJ_MAXVAL
dist2 = MJ_MAXVAL
pos1 = wp.vec3(0.0)
pos2 = wp.vec3(0.0)
nrm1 = wp.vec3(0.0)
nrm2 = wp.vec3(0.0)
cnt = 0
box_rotT = wp.transpose(box_rot)
for vi in range(3):
vert = wp.vec3(0.0)
if vi == 0:
vert = t1
elif vi == 1:
vert = t2
else:
vert = t3
diff = vert - box_pos
local = box_rotT @ diff
maxaxis = 0
maxval = wp.abs(local[0]) - box_size[0]
for j in range(1, 3):
val = wp.abs(local[j]) - box_size[j]
if val > maxval:
maxval = val
maxaxis = j
inside = True
for j in range(3):
if wp.abs(local[j]) > box_size[j] + tri_radius:
inside = False
if inside and cnt < 2:
nrm_local = wp.vec3(0.0)
if maxaxis == 0:
nrm_local = wp.vec3(wp.sign(local[0]), 0.0, 0.0)
elif maxaxis == 1:
nrm_local = wp.vec3(0.0, wp.sign(local[1]), 0.0)
else:
nrm_local = wp.vec3(0.0, 0.0, wp.sign(local[2]))
nrm_global = box_rot @ nrm_local
d = maxval - tri_radius
offset = tri_radius + d * 0.5
p = vert - nrm_global * offset
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm_global
else:
dist2 = d
pos2 = p
nrm2 = nrm_global
cnt += 1
for i in range(8):
if cnt >= 2:
break
vec = wp.vec3(
wp.where(i & 1, box_size[0], -box_size[0]),
wp.where(i & 2, box_size[1], -box_size[1]),
wp.where(i & 4, box_size[2], -box_size[2]),
)
corner = box_rot @ vec + box_pos
d, p, n = sphere_triangle(corner, 0.0, t1, t2, t3, tri_radius)
if d < MJ_MAXVAL:
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = n
elif cnt == 1:
dist2 = d
pos2 = p
nrm2 = n
cnt += 1
return (
wp.vec2(dist1, dist2),
mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
mat23f(nrm1[0], nrm1[1], nrm1[2], nrm2[0], nrm2[1], nrm2[2]),
)
@wp.func
def capsule_triangle(
capsule_pos: wp.vec3,
capsule_axis: wp.vec3,
capsule_radius: float,
capsule_half_length: float,
t1: wp.vec3,
t2: wp.vec3,
t3: wp.vec3,
tri_radius: float,
) -> Tuple[wp.vec2, mat23f, mat23f]:
"""Core contact geometry calculation for capsule-triangle collision.
Port of mjraw_CapsuleTriangle from engine_collision_primitive.c
Args:
capsule_pos: Center position of the capsule.
capsule_axis: Unit axis direction of the capsule.
capsule_radius: Radius of the capsule.
capsule_half_length: Half-length of the capsule cylinder.
t1: Triangle vertex positions.
t2: Triangle vertex positions.
t3: Triangle vertex positions.
tri_radius: Triangle (flex element) radius.
Returns:
- wp.vec2 of distances for up to 2 contacts (MJ_MAXVAL if no collision).
- mat23f of contact positions (2 x vec3).
- mat23f of contact normals (2 x vec3).
"""
dist1 = MJ_MAXVAL
dist2 = MJ_MAXVAL
pos1 = wp.vec3(0.0)
pos2 = wp.vec3(0.0)
nrm1 = wp.vec3(0.0)
nrm2 = wp.vec3(0.0)
cnt = 0
p1 = capsule_pos - capsule_axis * capsule_half_length
p2 = capsule_pos + capsule_axis * capsule_half_length
d, p, n = sphere_triangle(p1, capsule_radius, t1, t2, t3, tri_radius)
if d < MJ_MAXVAL:
dist1 = d
pos1 = p
nrm1 = n
cnt = 1
d, p, n = sphere_triangle(p2, capsule_radius, t1, t2, t3, tri_radius)
if d < MJ_MAXVAL and cnt < 2:
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = n
else:
dist2 = d
pos2 = p
nrm2 = n
cnt += 1
ab = p2 - p1
ab_len_sq = 4.0 * capsule_half_length * capsule_half_length
for vi in range(3):
if cnt >= 2:
break
vert = wp.vec3(0.0)
if vi == 0:
vert = t1
elif vi == 1:
vert = t2
else:
vert = t3
vec = vert - p1
t_param = wp.dot(vec, ab) / wp.max(MJ_MINVAL, ab_len_sq)
if t_param > MJ_MINVAL and t_param < 1.0 - MJ_MINVAL:
closest = p1 + ab * t_param
diff = vert - closest
dist_raw = wp.length(diff)
if dist_raw > MJ_MINVAL:
nrm = diff / dist_raw
d = dist_raw - capsule_radius - tri_radius
p = (closest + vert + nrm * (capsule_radius - tri_radius)) * 0.5
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm
else:
dist2 = d
pos2 = p
nrm2 = nrm
cnt += 1
return (
wp.vec2(dist1, dist2),
mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
mat23f(nrm1[0], nrm1[1], nrm1[2], nrm2[0], nrm2[1], nrm2[2]),
)
@wp.func
def cylinder_triangle(
cylinder_pos: wp.vec3,
cylinder_axis: wp.vec3,
cylinder_radius: float,
cylinder_half_height: float,
t1: wp.vec3,
t2: wp.vec3,
t3: wp.vec3,
tri_radius: float,
) -> Tuple[wp.vec2, mat23f, mat23f]:
"""Core contact geometry calculation for cylinder-triangle collision.
Args:
cylinder_pos: Center position of the cylinder.
cylinder_axis: Unit axis direction of the cylinder.
cylinder_radius: Radius of the cylinder.
cylinder_half_height: Half-height of the cylinder.
t1: Triangle vertex positions.
t2: Triangle vertex positions.
t3: Triangle vertex positions.
tri_radius: Triangle (flex element) radius.
Returns:
- wp.vec2 of distances for up to 2 contacts (MJ_MAXVAL if no collision).
- mat23f of contact positions (2 x vec3).
- mat23f of contact normals (2 x vec3).
"""
dist1 = MJ_MAXVAL
dist2 = MJ_MAXVAL
pos1 = wp.vec3(0.0)
pos2 = wp.vec3(0.0)
nrm1 = wp.vec3(0.0)
nrm2 = wp.vec3(0.0)
cnt = int(0)
p1 = cylinder_pos - cylinder_axis * cylinder_half_height
p2 = cylinder_pos + cylinder_axis * cylinder_half_height
ab = p2 - p1
ab_len_sq = 4.0 * cylinder_half_height * cylinder_half_height
for vi in range(3):
if cnt >= 2:
break
vert = wp.vec3(0.0)
if vi == 0:
vert = t1
elif vi == 1:
vert = t2
else:
vert = t3
vec = vert - p1
t_param = wp.dot(vec, ab) / wp.max(MJ_MINVAL, ab_len_sq)
if t_param > MJ_MINVAL and t_param < 1.0 - MJ_MINVAL:
closest = p1 + ab * t_param
diff = vert - closest
dist_raw = wp.length(diff)
if dist_raw < cylinder_radius + tri_radius:
if dist_raw > MJ_MINVAL:
nrm = diff / dist_raw
d = dist_raw - cylinder_radius - tri_radius
p = (closest + vert + nrm * (cylinder_radius - tri_radius)) * 0.5
else:
dist_to_side = cylinder_radius
dist_to_p2 = (1.0 - t_param) * wp.sqrt(ab_len_sq)
dist_to_p1 = t_param * wp.sqrt(ab_len_sq)
if dist_to_p2 < dist_to_side and dist_to_p2 < dist_to_p1:
nrm = cylinder_axis
d = -dist_to_p2 - tri_radius
p = vert
elif dist_to_p1 < dist_to_side:
nrm = -cylinder_axis
d = -dist_to_p1 - tri_radius
p = vert
else:
tri_normal = wp.normalize(wp.cross(t2 - t1, t3 - t1))
nrm = tri_normal
d = -cylinder_radius - tri_radius
p = closest
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm
else:
dist2 = d
pos2 = p
nrm2 = nrm
cnt += 1
elif t_param <= MJ_MINVAL:
diff = vert - p1
signed_dist = wp.dot(diff, cylinder_axis)
perp = diff - cylinder_axis * signed_dist
perp_len = wp.length(perp)
if perp_len < cylinder_radius:
d = -signed_dist - tri_radius
nrm = -cylinder_axis
p = vert - nrm * (tri_radius + d * 0.5)
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm
else:
dist2 = d
pos2 = p
nrm2 = nrm
cnt += 1
elif perp_len < cylinder_radius + tri_radius:
edge_dir = perp / perp_len
edge_point = p1 + edge_dir * cylinder_radius
diff_to_edge = vert - edge_point
dist_raw = wp.length(diff_to_edge)
if dist_raw > MJ_MINVAL:
nrm = diff_to_edge / dist_raw
d = dist_raw - tri_radius
p = vert - nrm * (tri_radius + d * 0.5)
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm
else:
dist2 = d
pos2 = p
nrm2 = nrm
cnt += 1
else:
diff = vert - p2
signed_dist = wp.dot(diff, cylinder_axis)
perp = diff - cylinder_axis * signed_dist
perp_len = wp.length(perp)
if perp_len < cylinder_radius:
d = signed_dist - tri_radius
nrm = cylinder_axis
p = vert - nrm * (tri_radius + d * 0.5)
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm
else:
dist2 = d
pos2 = p
nrm2 = nrm
cnt += 1
elif perp_len < cylinder_radius + tri_radius:
edge_dir = perp / perp_len
edge_point = p2 + edge_dir * cylinder_radius
diff_to_edge = vert - edge_point
dist_raw = wp.length(diff_to_edge)
if dist_raw > MJ_MINVAL:
nrm = diff_to_edge / dist_raw
d = dist_raw - tri_radius
p = vert - nrm * (tri_radius + d * 0.5)
if cnt == 0:
dist1 = d
pos1 = p
nrm1 = nrm
else:
dist2 = d
pos2 = p
nrm2 = nrm
cnt += 1
return (
wp.vec2(dist1, dist2),
mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
mat23f(nrm1[0], nrm1[1], nrm1[2], nrm2[0], nrm2[1], nrm2[2]),
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+222 -219
View File
@@ -15,6 +15,8 @@
from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
@@ -27,9 +29,9 @@ 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 vec8
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
from mujoco.mjx.third_party.mujoco_warp._src.types import vec_pluginattr
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
import warp as wp
wp.set_module_options({"enable_backward": False})
@@ -38,8 +40,8 @@ wp.set_module_options({"enable_backward": False})
class OptimizationParams:
rel_mat: wp.mat33
rel_pos: wp.vec3
attr1: wp.vec3
attr2: wp.vec3
attr1: vec_pluginattr
attr2: vec_pluginattr
@wp.struct
@@ -77,20 +79,24 @@ class MeshData:
@wp.func
def get_sdf_params(
# Model:
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8),
mesh_octadr: wp.array(dtype=int),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
# In:
g_type: int,
g_size: wp.vec3,
plugin_id: int,
mesh_id: int,
) -> Tuple[wp.vec3, int, VolumeData, MeshData]:
attributes = g_size
# Model:
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8),
mesh_octadr: wp.array(dtype=int),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=vec_pluginattr),
# In:
g_type: int,
g_size: wp.vec3,
plugin_id: int,
mesh_id: int,
) -> Tuple[vec_pluginattr, int, VolumeData, MeshData]:
# default attributes from geom size, first 3 values copied
attributes = vec_pluginattr()
attributes[0] = g_size[0]
attributes[1] = g_size[1]
attributes[2] = g_size[2]
plugin_index = -1
volume_data = VolumeData()
@@ -108,6 +114,16 @@ def get_sdf_params(
volume_data.oct_coeff = oct_coeff
volume_data.valid = True
elif g_type == GeomType.MESH and mesh_id != -1 and mesh_octadr[mesh_id] != -1:
octadr = mesh_octadr[mesh_id]
volume_data.center = oct_aabb[octadr, 0]
volume_data.half_size = oct_aabb[octadr, 1]
volume_data.root = octadr
volume_data.oct_aabb = oct_aabb
volume_data.oct_child = oct_child
volume_data.oct_coeff = oct_coeff
volume_data.valid = True
return attributes, plugin_index, volume_data, MeshData()
@@ -215,24 +231,28 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
@wp.func
def user_sdf(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float:
def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
"""User-defined SDF function.
Access attributes via attr[i] where i is the attribute index (0 to _NPLUGINATTR-1).
"""
wp.printf("ERROR: user_sdf function must be implemented by user code\n")
return 0.0
@wp.func
def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
def user_sdf_grad(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> wp.vec3:
"""User-defined SDF gradient function.
Access attributes via attr[i] where i is the attribute index (0 to _NPLUGINATTR-1).
"""
wp.printf("ERROR: user_sdf_grad function must be implemented by user code\n")
return wp.vec3(0.0)
@wp.func
def find_oct(
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
p: wp.vec3,
grad: bool,
root: int,
oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool, root: int
) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
stack = root
niter = int(100)
@@ -268,14 +288,14 @@ def find_oct(
# child indices are relative to root (mesh_octadr offset)
child0 = oct_child[node][0]
if (
child0 == -1
and oct_child[node][1] == -1
and oct_child[node][2] == -1
and oct_child[node][3] == -1
and oct_child[node][4] == -1
and oct_child[node][5] == -1
and oct_child[node][6] == -1
and oct_child[node][7] == -1
child0 == -1
and oct_child[node][1] == -1
and oct_child[node][2] == -1
and oct_child[node][3] == -1
and oct_child[node][4] == -1
and oct_child[node][5] == -1
and oct_child[node][6] == -1
and oct_child[node][7] == -1
):
for j in range(8):
if not grad:
@@ -342,13 +362,7 @@ def box_project(center: wp.vec3, half_size: wp.vec3, xyz: wp.vec3) -> Tuple[floa
@wp.func
def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float:
dist0, point = box_project(volume_data.center, volume_data.half_size, xyz)
node, weights = find_oct(
volume_data.oct_child,
volume_data.oct_aabb,
point,
grad=False,
root=volume_data.root,
)
node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=False, root=volume_data.root)
return dist0 + wp.dot(weights[0], volume_data.oct_coeff[node])
@@ -365,13 +379,7 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
grad_y = (sample_volume_sdf(xyz + dy, volume_data) - f) / h
grad_z = (sample_volume_sdf(xyz + dz, volume_data) - f) / h
return wp.vec3(grad_x, grad_y, grad_z)
node, weights = find_oct(
volume_data.oct_child,
volume_data.oct_aabb,
point,
grad=True,
root=volume_data.root,
)
node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=True, root=volume_data.root)
grad_x = wp.dot(weights[0], volume_data.oct_coeff[node])
grad_y = wp.dot(weights[1], volume_data.oct_coeff[node])
grad_z = wp.dot(weights[2], volume_data.oct_coeff[node])
@@ -379,15 +387,17 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
@wp.func
def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> float:
def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> float:
# extract first 3 elements as vec3 for primitive sdf functions
attr_vec3 = wp.vec3(attr[0], attr[1], attr[2])
if type == GeomType.PLANE:
return p[2]
elif type == GeomType.SPHERE:
return sphere(p, attr)
return sphere(p, attr_vec3)
elif type == GeomType.BOX:
return box(p, attr)
return box(p, attr_vec3)
elif type == GeomType.ELLIPSOID:
return ellipsoid(p, attr)
return ellipsoid(p, attr_vec3)
elif type == GeomType.MESH and mesh_data.valid:
mesh_data.pnt = p
mesh_data.vec = -wp.normalize(p)
@@ -425,21 +435,27 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
return sample_volume_sdf(p, volume_data)
else:
return user_sdf(p, attr, sdf_type)
elif type == GeomType.MESH and volume_data.valid:
return sample_volume_sdf(p, volume_data)
wp.printf("ERROR: SDF type not implemented\n")
return 0.0
@wp.func
def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> wp.vec3:
def sdf_grad(
type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData
) -> wp.vec3:
# extract first 3 elements as vec3 for primitive sdf functions
attr_vec3 = wp.vec3(attr[0], attr[1], attr[2])
if type == GeomType.PLANE:
grad = wp.vec3(0.0, 0.0, 1.0)
return grad
elif type == GeomType.SPHERE:
return grad_sphere(p)
elif type == GeomType.BOX:
return grad_box(p, attr)
return grad_box(p, attr_vec3)
elif type == GeomType.ELLIPSOID:
return grad_ellipsoid(p, attr)
return grad_ellipsoid(p, attr_vec3)
elif type == GeomType.MESH and mesh_data.valid:
mesh_data.pnt = p
mesh_data.vec = -wp.normalize(p)
@@ -466,6 +482,8 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V
return sample_volume_grad(p, volume_data)
else:
return user_sdf_grad(p, attr, sdf_type)
elif type == GeomType.MESH and volume_data.valid:
return sample_volume_grad(p, volume_data)
wp.printf("ERROR: SDF grad type not implemented\n")
return wp.vec3(0.0)
@@ -476,8 +494,8 @@ def clearance(
type1: int,
p1: wp.vec3,
p2: wp.vec3,
s1: wp.vec3,
s2: wp.vec3,
s1: vec_pluginattr,
s2: vec_pluginattr,
sdf_type1: int,
sdf_type2: int,
sfd_intersection: bool,
@@ -606,8 +624,8 @@ def gradient_descent(
# In:
type1: int,
x0_initial: wp.vec3,
attr1: wp.vec3,
attr2: wp.vec3,
attr1: vec_pluginattr,
attr2: vec_pluginattr,
pos1: wp.vec3,
rot1: wp.mat33,
pos2: wp.vec3,
@@ -645,76 +663,77 @@ def gradient_descent(
@wp.kernel
def _sdf_narrowphase(
# Model:
nmeshface: int,
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8),
geom_type: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_priority: wp.array(dtype=int),
geom_solmix: wp.array2d(dtype=float),
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_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
geom_gap: wp.array2d(dtype=float),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_faceadr: wp.array(dtype=int),
mesh_octadr: wp.array(dtype=int),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_face: wp.array(dtype=wp.vec3i),
mesh_graph: wp.array(dtype=int),
mesh_polynum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
mesh_polynormal: wp.array(dtype=wp.vec3),
mesh_polyvertadr: wp.array(dtype=int),
mesh_polyvertnum: wp.array(dtype=int),
mesh_polyvert: wp.array(dtype=int),
mesh_polymapadr: wp.array(dtype=int),
mesh_polymapnum: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
pair_dim: wp.array(dtype=int),
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
pair_solimp: wp.array2d(dtype=vec5),
pair_margin: wp.array2d(dtype=float),
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
geom_plugin_index: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
ncollision_in: wp.array(dtype=int),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
collision_worldid_in: wp.array(dtype=int),
sdf_initpoints: int,
sdf_iterations: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
# Model:
nmeshface: int,
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8),
geom_type: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_priority: wp.array(dtype=int),
geom_solmix: wp.array2d(dtype=float),
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_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
geom_gap: wp.array2d(dtype=float),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_faceadr: wp.array(dtype=int),
mesh_octadr: wp.array(dtype=int),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_face: wp.array(dtype=wp.vec3i),
mesh_graph: wp.array(dtype=int),
mesh_polynum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
mesh_polynormal: wp.array(dtype=wp.vec3),
mesh_polyvertadr: wp.array(dtype=int),
mesh_polyvertnum: wp.array(dtype=int),
mesh_polyvert: wp.array(dtype=int),
mesh_polymapadr: wp.array(dtype=int),
mesh_polymapnum: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
pair_dim: wp.array(dtype=int),
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
pair_solimp: wp.array2d(dtype=vec5),
pair_margin: wp.array2d(dtype=float),
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=vec_pluginattr),
geom_plugin_index: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
ncollision_in: wp.array(dtype=int),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
collision_worldid_in: wp.array(dtype=int),
sdf_initpoints: int,
sdf_iterations: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
i, contact_tid = wp.tid()
if i >= sdf_initpoints:
@@ -799,29 +818,11 @@ def _sdf_narrowphase(
rot1 = geom1.rot
attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params(
oct_child,
oct_aabb,
oct_coeff,
mesh_octadr,
plugin,
plugin_attr,
type1,
geom1.size,
g1_plugin,
geom_dataid[g1],
oct_child, oct_aabb, oct_coeff, mesh_octadr, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1]
)
attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params(
oct_child,
oct_aabb,
oct_coeff,
mesh_octadr,
plugin,
plugin_attr,
type2,
geom2.size,
g2_plugin,
geom_dataid[g2],
oct_child, oct_aabb, oct_coeff, mesh_octadr, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2]
)
mesh_data1.nmeshface = nmeshface
@@ -900,6 +901,7 @@ def _sdf_narrowphase(
contact_solimp_out,
contact_dim_out,
contact_geom_out,
contact_efc_address_out,
contact_worldid_out,
contact_type_out,
contact_geomcollisionid_out,
@@ -910,76 +912,77 @@ def _sdf_narrowphase(
@event_scope
def sdf_narrowphase(m: Model, d: Data, ctx: CollisionContext):
wp.launch(
_sdf_narrowphase,
dim=(m.opt.sdf_initpoints, d.naconmax),
inputs=[
m.nmeshface,
m.oct_child,
m.oct_aabb,
m.oct_coeff,
m.geom_type,
m.geom_condim,
m.geom_dataid,
m.geom_priority,
m.geom_solmix,
m.geom_solref,
m.geom_solimp,
m.geom_size,
m.geom_aabb,
m.geom_friction,
m.geom_margin,
m.geom_gap,
m.mesh_vertadr,
m.mesh_vertnum,
m.mesh_faceadr,
m.mesh_octadr,
m.mesh_graphadr,
m.mesh_vert,
m.mesh_face,
m.mesh_graph,
m.mesh_polynum,
m.mesh_polyadr,
m.mesh_polynormal,
m.mesh_polyvertadr,
m.mesh_polyvertnum,
m.mesh_polyvert,
m.mesh_polymapadr,
m.mesh_polymapnum,
m.mesh_polymap,
m.pair_dim,
m.pair_solref,
m.pair_solreffriction,
m.pair_solimp,
m.pair_margin,
m.pair_gap,
m.pair_friction,
m.plugin,
m.plugin_attr,
m.geom_plugin_index,
d.geom_xpos,
d.geom_xmat,
d.naconmax,
d.ncollision,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
m.opt.sdf_initpoints,
m.opt.sdf_iterations,
],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
],
_sdf_narrowphase,
dim=(m.opt.sdf_initpoints, d.naconmax),
inputs=[
m.nmeshface,
m.oct_child,
m.oct_aabb,
m.oct_coeff,
m.geom_type,
m.geom_condim,
m.geom_dataid,
m.geom_priority,
m.geom_solmix,
m.geom_solref,
m.geom_solimp,
m.geom_size,
m.geom_aabb,
m.geom_friction,
m.geom_margin,
m.geom_gap,
m.mesh_vertadr,
m.mesh_vertnum,
m.mesh_faceadr,
m.mesh_octadr,
m.mesh_graphadr,
m.mesh_vert,
m.mesh_face,
m.mesh_graph,
m.mesh_polynum,
m.mesh_polyadr,
m.mesh_polynormal,
m.mesh_polyvertadr,
m.mesh_polyvertnum,
m.mesh_polyvert,
m.mesh_polymapadr,
m.mesh_polymapnum,
m.mesh_polymap,
m.pair_dim,
m.pair_solref,
m.pair_solreffriction,
m.pair_solimp,
m.pair_margin,
m.pair_gap,
m.pair_friction,
m.plugin,
m.plugin_attr,
m.geom_plugin_index,
d.geom_xpos,
d.geom_xmat,
d.naconmax,
d.ncollision,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
m.opt.sdf_initpoints,
m.opt.sdf_iterations,
],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+1732 -1526
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+151 -44
View File
@@ -15,6 +15,7 @@
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
@@ -30,18 +31,24 @@ wp.set_module_options({"enable_backward": False})
@wp.kernel
def _qderiv_actuator_passive_vel(
# Model:
opt_timestep: wp.array(dtype=float),
actuator_dyntype: wp.array(dtype=int),
actuator_gaintype: wp.array(dtype=int),
actuator_biastype: wp.array(dtype=int),
actuator_actadr: wp.array(dtype=int),
actuator_actnum: wp.array(dtype=int),
actuator_forcelimited: wp.array(dtype=bool),
actuator_actlimited: wp.array(dtype=bool),
actuator_dynprm: wp.array2d(dtype=vec10f),
actuator_gainprm: wp.array2d(dtype=vec10f),
actuator_biasprm: wp.array2d(dtype=vec10f),
actuator_actearly: wp.array(dtype=bool),
actuator_forcerange: wp.array2d(dtype=wp.vec2),
actuator_actrange: wp.array2d(dtype=wp.vec2),
# Data in:
act_in: wp.array2d(dtype=float),
ctrl_in: wp.array2d(dtype=float),
act_dot_in: wp.array2d(dtype=float),
actuator_force_in: wp.array2d(dtype=float),
# Out:
vel_out: wp.array2d(dtype=float),
@@ -76,9 +83,24 @@ def _qderiv_actuator_passive_vel(
vel = float(bias)
if actuator_dyntype[actid] != DynType.NONE:
if gain != 0.0:
act_first = actuator_actadr[actid]
act_last = act_first + actuator_actnum[actid] - 1
vel += gain * act_in[worldid, act_last]
act_adr = actuator_actadr[actid] + actuator_actnum[actid] - 1
# use next activation if actearly is set (matching forward pass)
if actuator_actearly[actid]:
act = next_act(
opt_timestep[worldid % opt_timestep.shape[0]],
actuator_dyntype[actid],
actuator_dynprm[worldid % actuator_dynprm.shape[0], actid],
actuator_actrange[worldid % actuator_actrange.shape[0], actid],
act_in[worldid, act_adr],
act_dot_in[worldid, act_adr],
1.0,
actuator_actlimited[actid],
)
else:
act = act_in[worldid, act_adr]
vel += gain * act
else:
if gain != 0.0:
vel += gain * ctrl_in[worldid, actid]
@@ -95,21 +117,20 @@ def _nonzero_mask(x: float) -> float:
@wp.kernel
def _qderiv_actuator_passive_actuation_sparse(
# Model:
nu: int,
is_sparse: bool,
# Data in:
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
# In:
vel_in: wp.array2d(dtype=float),
qMi: wp.array(dtype=int),
qMj: wp.array(dtype=int),
# Out:
qDeriv_out: wp.array3d(dtype=float),
def _qderiv_actuator_passive_actuation_dense(
# Model:
nu: int,
# Data in:
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
# In:
vel_in: wp.array2d(dtype=float),
qMi: wp.array(dtype=int),
qMj: wp.array(dtype=int),
# Out:
qDeriv_out: wp.array3d(dtype=float),
):
worldid, elemid = wp.tid()
@@ -142,12 +163,63 @@ def _qderiv_actuator_passive_actuation_sparse(
qderiv_contrib += moment_i * moment_j * vel
if is_sparse:
qDeriv_out[worldid, 0, elemid] = qderiv_contrib
else:
qDeriv_out[worldid, dofiid, dofjid] = qderiv_contrib
if dofiid != dofjid:
qDeriv_out[worldid, dofjid, dofiid] = qderiv_contrib
qDeriv_out[worldid, dofiid, dofjid] = qderiv_contrib
if dofiid != dofjid:
qDeriv_out[worldid, dofjid, dofiid] = qderiv_contrib
@wp.kernel
def _qderiv_actuator_passive_actuation_sparse(
# Model:
M_rownnz: wp.array(dtype=int),
M_rowadr: wp.array(dtype=int),
# Data in:
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
# In:
vel_in: wp.array2d(dtype=float),
qMj: wp.array(dtype=int),
# Out:
qDeriv_out: wp.array3d(dtype=float),
):
worldid, actid = wp.tid()
vel = vel_in[worldid, actid]
if vel == 0.0:
return
rownnz = moment_rownnz_in[worldid, actid]
rowadr = moment_rowadr_in[worldid, actid]
for i in range(rownnz):
rowadri = rowadr + i
moment_i = actuator_moment_in[worldid, rowadri]
if moment_i == 0.0:
continue
dofi = moment_colind_in[worldid, rowadri]
for j in range(i + 1):
rowadrj = rowadr + j
moment_j = actuator_moment_in[worldid, rowadrj]
if moment_j == 0.0:
continue
dofj = moment_colind_in[worldid, rowadrj]
contrib = moment_i * moment_j * vel
# Search the corresponding elemid
# TODO: This could be precalculated for improved performance
row = dofi
col = dofj
row_startk = M_rowadr[row] - 1
row_nnz = M_rownnz[row]
for k in range(row_nnz):
row_startk += 1
if qMj[row_startk] == col:
wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
break
@wp.kernel
@@ -176,7 +248,7 @@ def _qderiv_actuator_passive(
else:
qderiv = qDeriv_in[worldid, dofiid, dofjid]
if not opt_disableflags & DisableBit.DAMPER and dofiid == dofjid:
if not (opt_disableflags & DisableBit.DAMPER) and dofiid == dofjid:
qderiv -= dof_damping[worldid % dof_damping.shape[0], dofiid]
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
@@ -196,10 +268,13 @@ def _qderiv_tendon_damping(
# Model:
ntendon: int,
opt_timestep: wp.array(dtype=float),
ten_J_rownnz: wp.array(dtype=int),
ten_J_rowadr: wp.array(dtype=int),
ten_J_colind: wp.array(dtype=int),
tendon_damping: wp.array2d(dtype=float),
is_sparse: bool,
# Data in:
ten_J_in: wp.array3d(dtype=float),
ten_J_in: wp.array2d(dtype=float),
# In:
qMi: wp.array(dtype=int),
qMj: wp.array(dtype=int),
@@ -213,7 +288,24 @@ def _qderiv_tendon_damping(
qderiv = float(0.0)
tendon_damping_id = worldid % tendon_damping.shape[0]
for tenid in range(ntendon):
qderiv -= ten_J_in[worldid, tenid, dofiid] * ten_J_in[worldid, tenid, dofjid] * tendon_damping[tendon_damping_id, tenid]
damping = tendon_damping[tendon_damping_id, tenid]
if damping == 0.0:
continue
rownnz = ten_J_rownnz[tenid]
rowadr = ten_J_rowadr[tenid]
Ji = float(0.0)
Jj = float(0.0)
for k in range(rownnz):
if Ji != 0.0 and Jj != 0.0:
break
sparseid = rowadr + k
colind = ten_J_colind[sparseid]
if colind == dofiid:
Ji = ten_J_in[worldid, sparseid]
if colind == dofjid:
Jj = ten_J_in[worldid, sparseid]
qderiv -= Ji * Jj * damping
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
@@ -242,43 +334,47 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
# TODO(team): only clear elements not set by _qderiv_actuator_passive
out.zero_()
if m.nu > 0 and not m.opt.disableflags & DisableBit.ACTUATION:
if m.nu > 0 and not (m.opt.disableflags & DisableBit.ACTUATION):
vel = wp.empty((d.nworld, m.nu), dtype=float)
wp.launch(
_qderiv_actuator_passive_vel,
dim=(d.nworld, m.nu),
inputs=[
m.opt.timestep,
m.actuator_dyntype,
m.actuator_gaintype,
m.actuator_biastype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_forcelimited,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_gainprm,
m.actuator_biasprm,
m.actuator_actearly,
m.actuator_forcerange,
m.actuator_actrange,
d.act,
d.ctrl,
d.act_dot,
d.actuator_force,
],
outputs=[vel],
)
wp.launch(
if m.is_sparse:
wp.launch(
_qderiv_actuator_passive_actuation_sparse,
dim=(d.nworld, qMi.size),
inputs=[
m.nu,
m.is_sparse,
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind,
d.actuator_moment,
vel,
qMi,
qMj,
],
dim=(d.nworld, m.nu),
inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
outputs=[out],
)
)
else:
wp.launch(
_qderiv_actuator_passive_actuation_dense,
dim=(d.nworld, qMi.size),
inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
outputs=[out],
)
wp.launch(
_qderiv_actuator_passive,
dim=(d.nworld, qMi.size),
@@ -298,11 +394,22 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
# TODO(team): directly utilize qM for these settings
wp.copy(out, d.qM)
if not m.opt.disableflags & DisableBit.DAMPER:
if not (m.opt.disableflags & DisableBit.DAMPER):
wp.launch(
_qderiv_tendon_damping,
dim=(d.nworld, qMi.size),
inputs=[m.ntendon, m.opt.timestep, m.tendon_damping, m.is_sparse, d.ten_J, qMi, qMj],
inputs=[
m.ntendon,
m.opt.timestep,
m.ten_J_rownnz,
m.ten_J_rowadr,
m.ten_J_colind,
m.tendon_damping,
m.is_sparse,
d.ten_J,
qMi,
qMj,
],
outputs=[out],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+202 -219
View File
@@ -15,6 +15,8 @@
from typing import Optional
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
from mujoco.mjx.third_party.mujoco_warp._src import constraint
from mujoco.mjx.third_party.mujoco_warp._src import derivative
@@ -25,7 +27,9 @@ from mujoco.mjx.third_party.mujoco_warp._src import sensor
from mujoco.mjx.third_party.mujoco_warp._src import smooth
from mujoco.mjx.third_party.mujoco_warp._src import solver
from mujoco.mjx.third_party.mujoco_warp._src import util_misc
from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
@@ -34,14 +38,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
from mujoco.mjx.third_party.mujoco_warp._src.types import 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
import warp as wp
wp.set_module_options({"enable_backward": False})
@@ -127,55 +129,24 @@ def _next_velocity(
qvel_out[worldid, dofid] = qvel_in[worldid, dofid] + qacc_scale_in * qacc_in[worldid, dofid] * timestep
# TODO(team): kernel analyzer array slice?
@wp.func
def _next_act(
# Model:
opt_timestep: float, # kernel_analyzer: ignore
actuator_dyntype: int, # kernel_analyzer: ignore
actuator_dynprm: vec10f, # kernel_analyzer: ignore
actuator_actrange: wp.vec2, # kernel_analyzer: ignore
# Data In:
act_in: float, # kernel_analyzer: ignore
act_dot_in: float, # kernel_analyzer: ignore
# In:
act_dot_scale: float,
clamp: bool,
) -> float:
# advance actuation
if actuator_dyntype == DynType.FILTEREXACT:
tau = wp.max(MJ_MINVAL, actuator_dynprm[0])
act = act_in + act_dot_scale * act_dot_in * tau * (1.0 - wp.exp(-opt_timestep / tau))
elif actuator_dyntype == DynType.USER:
return act_in
else:
act = act_in + act_dot_scale * act_dot_in * opt_timestep
# clamp to actrange
if clamp:
act = wp.clamp(act, actuator_actrange[0], actuator_actrange[1])
return act
@wp.kernel
def _next_activation(
# Model:
opt_timestep: wp.array(dtype=float),
actuator_dyntype: wp.array(dtype=int),
actuator_actadr: wp.array(dtype=int),
actuator_actnum: wp.array(dtype=int),
actuator_actlimited: wp.array(dtype=bool),
actuator_dynprm: wp.array2d(dtype=vec10f),
actuator_actrange: wp.array2d(dtype=wp.vec2),
# Data in:
act_in: wp.array2d(dtype=float),
act_dot_in: wp.array2d(dtype=float),
# In:
act_dot_scale: float,
limit: bool,
# Data out:
act_out: wp.array2d(dtype=float),
# Model:
opt_timestep: wp.array(dtype=float),
actuator_dyntype: wp.array(dtype=int),
actuator_actadr: wp.array(dtype=int),
actuator_actnum: wp.array(dtype=int),
actuator_actlimited: wp.array(dtype=bool),
actuator_dynprm: wp.array2d(dtype=vec10f),
actuator_actrange: wp.array2d(dtype=wp.vec2),
# Data in:
act_in: wp.array2d(dtype=float),
act_dot_in: wp.array2d(dtype=float),
# In:
act_dot_scale: float,
limit: bool,
# Data out:
act_out: wp.array2d(dtype=float),
):
worldid, uid = wp.tid()
opt_timestep_id = worldid % opt_timestep.shape[0]
@@ -184,15 +155,15 @@ def _next_activation(
actadr = actuator_actadr[uid]
actnum = actuator_actnum[uid]
for j in range(actadr, actadr + actnum):
act = _next_act(
opt_timestep[opt_timestep_id],
actuator_dyntype[uid],
actuator_dynprm[actuator_dynprm_id, uid],
actuator_actrange[actuator_actrange_id, uid],
act_in[worldid, j],
act_dot_in[worldid, j],
act_dot_scale,
limit and actuator_actlimited[uid],
act = next_act(
opt_timestep[opt_timestep_id],
actuator_dyntype[uid],
actuator_dynprm[actuator_dynprm_id, uid],
actuator_actrange[actuator_actrange_id, uid],
act_in[worldid, j],
act_dot_in[worldid, j],
act_dot_scale,
limit and actuator_actlimited[uid],
)
act_out[worldid, j] = act
@@ -201,12 +172,16 @@ def _next_activation(
def _next_time(
# Model:
opt_timestep: wp.array(dtype=float),
is_sparse: bool,
# Data in:
nefc_in: wp.array(dtype=int),
time_in: wp.array(dtype=float),
efc_J_rownnz_in: wp.array2d(dtype=int),
efc_J_rowadr_in: wp.array2d(dtype=int),
nworld_in: int,
naconmax_in: int,
njmax_in: int,
njmax_nnz_in: int,
nacon_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int),
# Data out:
@@ -218,6 +193,11 @@ def _next_time(
if nefc > njmax_in:
wp.printf("nefc overflow - please increase njmax to %u\n", nefc)
elif nefc > 0 and is_sparse:
efcid = wp.min(nefc, njmax_in) - 1
efc_nnz = efc_J_rowadr_in[worldid, efcid] + efc_J_rownnz_in[worldid, efcid]
if efc_nnz > njmax_nnz_in:
wp.printf("njmax_nnz overflow - please increase njmax_nnz to %u\n", efc_nnz)
if worldid == 0:
ncollision = ncollision_in[0]
@@ -236,22 +216,22 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
# advance activations
wp.launch(
_next_activation,
dim=(d.nworld, m.nu),
inputs=[
m.opt.timestep,
m.actuator_dyntype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_actrange,
d.act,
d.act_dot,
1.0,
True,
],
outputs=[d.act],
_next_activation,
dim=(d.nworld, m.nu),
inputs=[
m.opt.timestep,
m.actuator_dyntype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_actrange,
d.act,
d.act_dot,
1.0,
True,
],
outputs=[d.act],
)
wp.launch(
@@ -274,7 +254,20 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
wp.launch(
_next_time,
dim=d.nworld,
inputs=[m.opt.timestep, d.nefc, d.time, d.nworld, d.naconmax, d.njmax, d.nacon, d.ncollision],
inputs=[
m.opt.timestep,
m.is_sparse,
d.nefc,
d.time,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.nworld,
d.naconmax,
d.njmax,
d.njmax_nnz,
d.nacon,
d.ncollision,
],
outputs=[d.time],
)
@@ -294,9 +287,7 @@ def _euler_damp_qfrc_sparse(
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
adr = dof_Madr[tid]
qM_integration_out[worldid, 0, adr] += (
timestep * dof_damping[worldid % dof_damping.shape[0], tid]
)
qM_integration_out[worldid, 0, adr] += timestep * dof_damping[worldid % dof_damping.shape[0], tid]
@cache_kernel
@@ -336,7 +327,7 @@ def _tile_euler_dense(tile: TileSet):
def euler(m: Model, d: Data):
"""Euler integrator, semi-implicit in velocity."""
# integrate damping implicitly
if not m.opt.disableflags & (DisableBit.EULERDAMP | DisableBit.DAMPER):
if not (m.opt.disableflags & (DisableBit.EULERDAMP | DisableBit.DAMPER)):
qacc = wp.empty((d.nworld, m.nv), dtype=float)
if m.is_sparse:
qM = wp.clone(d.qM)
@@ -390,22 +381,22 @@ def _rk_perturb_state(
# activation
if m.na and act_t0 is not None:
wp.launch(
_next_activation,
dim=(d.nworld, m.nu),
inputs=[
m.opt.timestep,
m.actuator_dyntype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_actrange,
act_t0,
d.act_dot,
scale,
False,
],
outputs=[d.act],
_next_activation,
dim=(d.nworld, m.nu),
inputs=[
m.opt.timestep,
m.actuator_dyntype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_actrange,
act_t0,
d.act_dot,
scale,
False,
],
outputs=[d.act],
)
@@ -548,14 +539,14 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
@wp.kernel
def _actuator_velocity(
# Data in:
qvel_in: wp.array2d(dtype=float),
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
# Data out:
actuator_velocity_out: wp.array2d(dtype=float),
# Data in:
qvel_in: wp.array2d(dtype=float),
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
# Data out:
actuator_velocity_out: wp.array2d(dtype=float),
):
worldid, actid = wp.tid()
@@ -571,50 +562,49 @@ def _actuator_velocity(
actuator_velocity_out[worldid, actid] = vel
@cache_kernel
def _tendon_velocity(nv: int):
@wp.kernel(module="unique", enable_backward=False)
def tendon_velocity(
# Data in:
qvel_in: wp.array2d(dtype=float),
ten_J_in: wp.array3d(dtype=float),
# Data out:
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(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.kernel
def _tendon_velocity(
# Model:
ten_J_rownnz: wp.array(dtype=int),
ten_J_rowadr: wp.array(dtype=int),
ten_J_colind: wp.array(dtype=int),
# Data in:
qvel_in: wp.array2d(dtype=float),
ten_J_in: wp.array2d(dtype=float),
# Data out:
ten_velocity_out: wp.array2d(dtype=float),
):
worldid, tenid = wp.tid()
return tendon_velocity
velocity = float(0.0)
rownnz = ten_J_rownnz[tenid]
rowadr = ten_J_rowadr[tenid]
for i in range(rownnz):
sparseid = rowadr + i
J = ten_J_in[worldid, sparseid]
if J != 0.0:
colind = ten_J_colind[sparseid]
velocity += J * qvel_in[worldid, colind]
ten_velocity_out[worldid, tenid] = velocity
@event_scope
def fwd_velocity(m: Model, d: Data):
"""Velocity-dependent computations."""
wp.launch_tiled(
_actuator_velocity,
dim=(d.nworld, m.nu),
inputs=[
d.qvel,
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind,
d.actuator_moment,
],
outputs=[d.actuator_velocity],
block_dim=m.block_dim.actuator_velocity,
wp.launch(
_actuator_velocity,
dim=(d.nworld, m.nu),
inputs=[d.qvel, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment],
outputs=[d.actuator_velocity],
block_dim=m.block_dim.actuator_velocity,
)
# TODO(team): sparse version
wp.launch_tiled(
_tendon_velocity(m.nv),
wp.launch(
_tendon_velocity,
dim=(d.nworld, m.ntendon),
inputs=[d.qvel, d.ten_J],
inputs=[m.ten_J_rownnz, m.ten_J_rowadr, m.ten_J_colind, d.qvel, d.ten_J],
outputs=[d.ten_velocity],
block_dim=m.block_dim.tendon_velocity,
)
smooth.com_vel(m, d)
@@ -625,36 +615,36 @@ def fwd_velocity(m: Model, d: Data):
@wp.kernel
def _actuator_force(
# Model:
na: int,
opt_timestep: wp.array(dtype=float),
actuator_dyntype: wp.array(dtype=int),
actuator_gaintype: wp.array(dtype=int),
actuator_biastype: wp.array(dtype=int),
actuator_actadr: wp.array(dtype=int),
actuator_actnum: wp.array(dtype=int),
actuator_ctrllimited: wp.array(dtype=bool),
actuator_forcelimited: wp.array(dtype=bool),
actuator_actlimited: wp.array(dtype=bool),
actuator_dynprm: wp.array2d(dtype=vec10f),
actuator_gainprm: wp.array2d(dtype=vec10f),
actuator_biasprm: wp.array2d(dtype=vec10f),
actuator_actearly: wp.array(dtype=bool),
actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
actuator_forcerange: wp.array2d(dtype=wp.vec2),
actuator_actrange: wp.array2d(dtype=wp.vec2),
actuator_acc0: wp.array2d(dtype=float),
actuator_lengthrange: wp.array2d(dtype=wp.vec2),
# Data in:
act_in: wp.array2d(dtype=float),
ctrl_in: wp.array2d(dtype=float),
actuator_length_in: wp.array2d(dtype=float),
actuator_velocity_in: wp.array2d(dtype=float),
# In:
dsbl_clampctrl: int,
# Data out:
act_dot_out: wp.array2d(dtype=float),
actuator_force_out: wp.array2d(dtype=float),
# Model:
na: int,
opt_timestep: wp.array(dtype=float),
actuator_dyntype: wp.array(dtype=int),
actuator_gaintype: wp.array(dtype=int),
actuator_biastype: wp.array(dtype=int),
actuator_actadr: wp.array(dtype=int),
actuator_actnum: wp.array(dtype=int),
actuator_ctrllimited: wp.array(dtype=bool),
actuator_forcelimited: wp.array(dtype=bool),
actuator_actlimited: wp.array(dtype=bool),
actuator_dynprm: wp.array2d(dtype=vec10f),
actuator_gainprm: wp.array2d(dtype=vec10f),
actuator_biasprm: wp.array2d(dtype=vec10f),
actuator_actearly: wp.array(dtype=bool),
actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
actuator_forcerange: wp.array2d(dtype=wp.vec2),
actuator_actrange: wp.array2d(dtype=wp.vec2),
actuator_acc0: wp.array2d(dtype=float),
actuator_lengthrange: wp.array2d(dtype=wp.vec2),
# Data in:
act_in: wp.array2d(dtype=float),
ctrl_in: wp.array2d(dtype=float),
actuator_length_in: wp.array2d(dtype=float),
actuator_velocity_in: wp.array2d(dtype=float),
# In:
dsbl_clampctrl: int,
# Data out:
act_dot_out: wp.array2d(dtype=float),
actuator_force_out: wp.array2d(dtype=float),
):
worldid, uid = wp.tid()
@@ -693,7 +683,7 @@ def _actuator_force(
if dyntype == DynType.INTEGRATOR or dyntype == DynType.NONE:
act = act_in[worldid, act_last]
ctrl_act = _next_act(
ctrl_act = next_act(
opt_timestep[worldid % opt_timestep.shape[0]],
dyntype,
dynprm,
@@ -720,9 +710,7 @@ def _actuator_force(
gain = gainprm[0] + gainprm[1] * length + gainprm[2] * velocity
elif gaintype == GainType.MUSCLE:
acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid]
lengthrange = actuator_lengthrange[
worldid % actuator_lengthrange.shape[0], uid
]
lengthrange = actuator_lengthrange[worldid % actuator_lengthrange.shape[0], uid]
gain = util_misc.muscle_gain(length, velocity, lengthrange, acc0, gainprm)
# GainType.USER: gain stays 0, modified by act_gain_callback
@@ -735,9 +723,7 @@ def _actuator_force(
bias = biasprm[0] + biasprm[1] * length + biasprm[2] * velocity
elif biastype == BiasType.MUSCLE:
acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid]
lengthrange = actuator_lengthrange[
worldid % actuator_lengthrange.shape[0], uid
]
lengthrange = actuator_lengthrange[worldid % actuator_lengthrange.shape[0], uid]
bias = util_misc.muscle_bias(length, lengthrange, acc0, biasprm)
force = gain * ctrl_act + bias
@@ -795,14 +781,14 @@ def _tendon_actuator_force_clamp(
@wp.kernel
def _qfrc_actuator(
# Data in:
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
actuator_force_in: wp.array2d(dtype=float),
# Data out:
qfrc_actuator_out: wp.array2d(dtype=float),
# Data in:
moment_rownnz_in: wp.array2d(dtype=int),
moment_rowadr_in: wp.array2d(dtype=int),
moment_colind_in: wp.array2d(dtype=int),
actuator_moment_in: wp.array2d(dtype=float),
actuator_force_in: wp.array2d(dtype=float),
# Data out:
qfrc_actuator_out: wp.array2d(dtype=float),
):
worldid, actid = wp.tid()
@@ -812,26 +798,23 @@ def _qfrc_actuator(
for i in range(rownnz):
sparseid = rowadr + i
colind = moment_colind_in[worldid, sparseid]
qfrc = (
actuator_moment_in[worldid, sparseid]
* actuator_force_in[worldid, actid]
)
qfrc = actuator_moment_in[worldid, sparseid] * actuator_force_in[worldid, actid]
wp.atomic_add(qfrc_actuator_out[worldid], colind, qfrc)
@wp.kernel
def _qfrc_actuator_gravcomp_limits(
# Model:
ngravcomp: int,
jnt_actfrclimited: wp.array(dtype=bool),
jnt_actgravcomp: wp.array(dtype=int),
jnt_actfrcrange: wp.array2d(dtype=wp.vec2),
dof_jntid: wp.array(dtype=int),
# Data in:
qfrc_gravcomp_in: wp.array2d(dtype=float),
qfrc_actuator_in: wp.array2d(dtype=float),
# Data out:
qfrc_actuator_out: wp.array2d(dtype=float),
# Model:
ngravcomp: int,
jnt_actfrclimited: wp.array(dtype=bool),
jnt_actgravcomp: wp.array(dtype=int),
jnt_actfrcrange: wp.array2d(dtype=wp.vec2),
dof_jntid: wp.array(dtype=int),
# Data in:
qfrc_gravcomp_in: wp.array2d(dtype=float),
qfrc_actuator_in: wp.array2d(dtype=float),
# Data out:
qfrc_actuator_out: wp.array2d(dtype=float),
):
worldid, dofid = wp.tid()
jntid = dof_jntid[dofid]
@@ -917,30 +900,30 @@ def fwd_actuation(m: Model, d: Data):
# TODO(team): optimize performance
d.qfrc_actuator.zero_()
wp.launch(
_qfrc_actuator,
dim=(d.nworld, m.nu),
inputs=[
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind,
d.actuator_moment,
d.actuator_force,
],
outputs=[d.qfrc_actuator],
_qfrc_actuator,
dim=(d.nworld, m.nu),
inputs=[
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind,
d.actuator_moment,
d.actuator_force,
],
outputs=[d.qfrc_actuator],
)
wp.launch(
_qfrc_actuator_gravcomp_limits,
dim=(d.nworld, m.nv),
inputs=[
m.ngravcomp,
m.jnt_actfrclimited,
m.jnt_actgravcomp,
m.jnt_actfrcrange,
m.dof_jntid,
d.qfrc_gravcomp,
d.qfrc_actuator,
],
outputs=[d.qfrc_actuator],
_qfrc_actuator_gravcomp_limits,
dim=(d.nworld, m.nv),
inputs=[
m.ngravcomp,
m.jnt_actfrclimited,
m.jnt_actgravcomp,
m.jnt_actfrcrange,
m.dof_jntid,
d.qfrc_gravcomp,
d.qfrc_actuator,
],
outputs=[d.qfrc_actuator],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+501 -373
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+17 -16
View File
@@ -13,12 +13,13 @@
# limitations under the License.
# ==============================================================================
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import types
from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
import warp as wp
@wp.kernel
@@ -180,17 +181,17 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)):
@wp.kernel
def _flood_fill(
# Model:
ntree: int,
# In:
tree_tree_in: wp.array3d(dtype=int),
labels_in: wp.array2d(dtype=int),
stack_in: wp.array2d(dtype=int),
# Data out:
nisland_out: wp.array(dtype=int),
tree_island_out: wp.array2d(dtype=int),
# Out:
stack_out: wp.array2d(dtype=int),
# Model:
ntree: int,
# In:
tree_tree_in: wp.array3d(dtype=int),
labels_in: wp.array2d(dtype=int),
stack_in: wp.array2d(dtype=int),
# Data out:
nisland_out: wp.array(dtype=int),
tree_island_out: wp.array2d(dtype=int),
# Out:
stack_out: wp.array2d(dtype=int),
):
"""DFS flood fill to discover islands using tree_tree matrix."""
worldid = wp.tid()
@@ -257,8 +258,8 @@ def island(m: types.Model, d: types.Data):
stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
wp.launch(
_flood_fill,
dim=d.nworld,
inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
outputs=[d.nisland, d.tree_island, stack_scratch],
_flood_fill,
dim=d.nworld,
inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
outputs=[d.nisland, d.tree_island, stack_scratch],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/math.py
Vendored
+29
View File
@@ -83,6 +83,35 @@ def quat_to_mat(quat: wp.quat) -> wp.mat33:
)
@wp.func
def quat_z2vec(vec: wp.vec3) -> wp.quat:
"""Compute quaternion performing rotation from z-axis to given vector."""
quat = wp.quat(0.0, 0.0, 0.0, 1.0)
# normalize vector; if too small, no rotation
norm = wp.length(vec)
if norm < types.MJ_MINVAL:
return quat
vec = vec / norm
axis = wp.vec3(-vec[1], vec[0], 0.0)
a = wp.length(axis)
# almost parallel
if a < types.MJ_MINVAL:
# opposite: 180 deg rotation around x axis
if vec[2] < 0.0:
quat = wp.quat(1.0, 0.0, 0.0, 0.0)
return quat
# make quaternion from angle and axis
axis = axis / a
angle = wp.atan2(a, vec[2])
quat = axis_angle_to_quat(axis, angle)
return quat
@wp.func
def quat_inv(quat: wp.quat) -> wp.quat:
return wp.quat(quat[0], -quat[1], -quat[2], -quat[3])
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+81 -60
View File
@@ -89,8 +89,8 @@ def _spring_damper_dof_passive(
stiffness = jnt_stiffness[worldid % jnt_stiffness.shape[0], jntid]
damping = dof_damping[worldid % dof_damping.shape[0], dofid]
has_stiffness = stiffness != 0.0 and not opt_disableflags & DisableBit.SPRING
has_damping = damping != 0.0 and not opt_disableflags & DisableBit.DAMPER
has_stiffness = stiffness != 0.0 and not (opt_disableflags & DisableBit.SPRING)
has_damping = damping != 0.0 and not (opt_disableflags & DisableBit.DAMPER)
if not has_stiffness:
qfrc_spring_out[worldid, dofid] = 0.0
@@ -182,11 +182,14 @@ def _spring_damper_dof_passive(
@wp.kernel
def _spring_damper_tendon_passive(
# Model:
ten_J_rownnz: wp.array(dtype=int),
ten_J_rowadr: wp.array(dtype=int),
ten_J_colind: wp.array(dtype=int),
tendon_stiffness: wp.array2d(dtype=float),
tendon_damping: wp.array2d(dtype=float),
tendon_lengthspring: wp.array2d(dtype=wp.vec2),
# Data in:
ten_J_in: wp.array3d(dtype=float),
ten_J_in: wp.array2d(dtype=float),
ten_length_in: wp.array2d(dtype=float),
ten_velocity_in: wp.array2d(dtype=float),
# In:
@@ -196,7 +199,7 @@ def _spring_damper_tendon_passive(
qfrc_spring_out: wp.array2d(dtype=float),
qfrc_damper_out: wp.array2d(dtype=float),
):
worldid, tenid, dofid = wp.tid()
worldid, tenid, dofid_sparse = wp.tid()
stiffness = tendon_stiffness[worldid % tendon_stiffness.shape[0], tenid]
damping = tendon_damping[worldid % tendon_damping.shape[0], tenid]
@@ -207,7 +210,13 @@ def _spring_damper_tendon_passive(
if not has_stiffness and not has_damping:
return
J = ten_J_in[worldid, tenid, dofid]
rownnz = ten_J_rownnz[tenid]
if dofid_sparse >= rownnz:
return
rowadr = ten_J_rowadr[tenid]
sparseid = rowadr + dofid_sparse
J = ten_J_in[worldid, sparseid]
dofid = ten_J_colind[sparseid]
if has_stiffness:
# compute spring force along tendon
@@ -265,28 +274,28 @@ def _gravity_force(
@wp.kernel
def _fluid_force(
# Model:
opt_wind: wp.array(dtype=wp.vec3),
opt_density: wp.array(dtype=float),
opt_viscosity: wp.array(dtype=float),
body_rootid: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int),
body_geomadr: wp.array(dtype=int),
body_mass: wp.array2d(dtype=float),
body_inertia: wp.array2d(dtype=wp.vec3),
geom_type: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_fluid: wp.array2d(dtype=float),
body_fluid_ellipsoid: wp.array(dtype=bool),
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# Out:
fluid_applied_out: wp.array2d(dtype=wp.spatial_vector),
# Model:
opt_wind: wp.array(dtype=wp.vec3),
opt_density: wp.array(dtype=float),
opt_viscosity: wp.array(dtype=float),
body_rootid: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int),
body_geomadr: wp.array(dtype=int),
body_mass: wp.array2d(dtype=float),
body_inertia: wp.array2d(dtype=wp.vec3),
geom_type: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_fluid: wp.array2d(dtype=float),
body_fluid_ellipsoid: wp.array(dtype=bool),
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# Out:
fluid_applied_out: wp.array2d(dtype=wp.spatial_vector),
):
"""Computes body-space fluid forces for both inertia-box and ellipsoid models."""
worldid, bodyid = wp.tid()
@@ -495,29 +504,29 @@ def _fluid(m: Model, d: Data):
fluid_applied = wp.empty((d.nworld, m.nbody), dtype=wp.spatial_vector)
wp.launch(
_fluid_force,
dim=(d.nworld, m.nbody),
inputs=[
m.opt.wind,
m.opt.density,
m.opt.viscosity,
m.body_rootid,
m.body_geomnum,
m.body_geomadr,
m.body_mass,
m.body_inertia,
m.geom_type,
m.geom_size,
m.geom_fluid,
m.body_fluid_ellipsoid,
d.xipos,
d.ximat,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
],
outputs=[fluid_applied],
_fluid_force,
dim=(d.nworld, m.nbody),
inputs=[
m.opt.wind,
m.opt.density,
m.opt.viscosity,
m.body_rootid,
m.body_geomnum,
m.body_geomadr,
m.body_mass,
m.body_inertia,
m.geom_type,
m.geom_size,
m.geom_fluid,
m.body_fluid_ellipsoid,
d.xipos,
d.ximat,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
],
outputs=[fluid_applied],
)
support.apply_ft(m, d, fluid_applied, d.qfrc_fluid, False)
@@ -565,6 +574,7 @@ def _flex_elasticity(
flex_edgeadr: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elemedgeadr: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
@@ -590,32 +600,39 @@ def _flex_elasticity(
f = i
break
local_elemid = elemid - flex_elemadr[f]
dim = flex_dim[f]
nvert = dim + 1
nedge = nvert * (nvert - 1) / 2
edges = wp.where(
dim == 3,
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),
dim == 1,
wp.matrix(0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int),
wp.where(
dim == 3,
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
elem_data_adr = flex_elemdataadr[f] + local_elemid * (dim + 1)
vbase = flex_vertadr[f]
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]]
xpos0 = flexvert_xpos_in[worldid, vert0]
xpos1 = flexvert_xpos_in[worldid, vert1]
vert0 = flex_elem[elem_data_adr + edges[e, 0]]
vert1 = flex_elem[elem_data_adr + edges[e, 1]]
xpos0 = flexvert_xpos_in[worldid, vbase + vert0]
xpos1 = flexvert_xpos_in[worldid, vbase + vert1]
for i in range(3):
gradient[e, 0 + i] = xpos0[i] - xpos1[i]
gradient[e, 3 + i] = xpos1[i] - xpos0[i]
elongation = wp.spatial_vectorf(0.0)
for e in range(nedge):
idx = flex_elemedge[elemid * nedge + e]
idx = flex_elemedge[flex_elemedgeadr[f] + local_elemid * nedge + e]
vel = flexedge_velocity_in[worldid, flex_edgeadr[f] + idx]
deformed = flexedge_length_in[worldid, flex_edgeadr[f] + idx]
reference = flexedge_length0[flex_edgeadr[f] + idx]
@@ -638,7 +655,7 @@ def _flex_elasticity(
force[edges[ed2, i], x] -= elongation[ed1] * gradient[ed2, 3 * i + x] * metric[ed1, ed2]
for v in range(nvert):
vert = flex_elem[(dim + 1) * elemid + v]
vert = flex_elem[elem_data_adr + v]
bodyid = flex_vertbodyid[flex_vertadr[f] + vert]
for x in range(3):
wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[v, x])
@@ -742,8 +759,11 @@ def passive(m: Model, d: Data):
if m.ntendon:
wp.launch(
_spring_damper_tendon_passive,
dim=(d.nworld, m.ntendon, m.nv),
dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
inputs=[
m.ten_J_rownnz,
m.ten_J_rowadr,
m.ten_J_colind,
m.tendon_stiffness,
m.tendon_damping,
m.tendon_lengthspring,
@@ -772,6 +792,7 @@ def passive(m: Model, d: Data):
m.flex_edgeadr,
m.flex_elemadr,
m.flex_elemnum,
m.flex_elemdataadr,
m.flex_elemedgeadr,
m.flex_vertbodyid,
m.flex_elem,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+20 -2
View File
@@ -752,7 +752,8 @@ def ray_mesh_with_bvh_anyhit(
@wp.func
def ray_flex_with_bvh(
# In:
bvh_id: wp.uint64,
flex_bvh_id: wp.array(dtype=wp.uint64),
flexid: int,
group_root: int,
pnt: wp.vec3,
vec: wp.vec3,
@@ -769,7 +770,7 @@ def ray_flex_with_bvh(
n = wp.vec3(0.0, 0.0, 0.0)
f = int(-1)
hit = wp.mesh_query_ray(bvh_id, pnt, vec, max_t, t, u, v, sign, n, f, group_root)
hit = wp.mesh_query_ray(flex_bvh_id[flexid], pnt, vec, max_t, t, u, v, sign, n, f, group_root)
if hit:
return t, n, u, v, f
@@ -777,6 +778,23 @@ def ray_flex_with_bvh(
return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1
@wp.func
def ray_flex_with_bvh_anyhit(
# In:
flex_bvh_id: wp.array(dtype=wp.uint64),
flexid: int,
group_root: int,
pnt: wp.vec3,
vec: wp.vec3,
max_t: float,
) -> bool:
"""Returns True if there is any hit for ray flex intersections.
Requires wp.Mesh be constructed and their ids to be passed. Flex are already in world space.
"""
return wp.mesh_query_ray_anyhit(flex_bvh_id[flexid], pnt, vec, max_t, group_root)
@wp.func
def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> Tuple[float, wp.vec3]:
"""Returns distance along ray to intersection with geom and normal at intersection point.
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+258 -126
View File
@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_capsule
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_cylinder
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_ellipsoid
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_flex_with_bvh
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_flex_with_bvh_anyhit
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh_with_bvh
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh_with_bvh_anyhit
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_plane
@@ -39,10 +40,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
from mujoco.mjx.third_party.mujoco_warp._src.types import TEXTURE_DTYPE
@wp.func
def sample_texture(
# Model:
@@ -51,7 +48,7 @@ def sample_texture(
# In:
geom_id: int,
tex_repeat: wp.vec2,
tex: TEXTURE_DTYPE,
tex: wp.Texture2D,
pos: wp.vec3,
rot: wp.mat33,
mesh_facetexcoord: wp.array(dtype=wp.vec3i),
@@ -94,17 +91,26 @@ def cast_ray(
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
flex_vertadr: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_radius: wp.array(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
bvh_id: wp.uint64,
group_root: int,
world_id: int,
worldid: int,
bvh_ngeom: int,
flex_bvh_ngeom: int,
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
flex_geom_flexid: wp.array(dtype=int),
flex_geom_edgeid: wp.array(dtype=int),
flex_bvh_id: wp.array(dtype=wp.uint64),
flex_group_root: wp.array2d(dtype=int),
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
) -> Tuple[int, float, wp.vec3, float, float, int, int]:
@@ -118,91 +124,127 @@ def cast_ray(
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
bounds_nr = int(0)
ngeom = bvh_ngeom + flex_bvh_ngeom
while wp.bvh_query_next(query, bounds_nr, dist):
gi_global = bounds_nr
gi_bvh_local = gi_global - (world_id * bvh_ngeom)
gi = enabled_geom_ids[gi_bvh_local]
local_id = gi_global - (worldid * ngeom)
d = float(-1.0)
hit_mesh_id = int(-1)
u = float(0.0)
v = float(0.0)
f = int(-1)
n = wp.vec3(0.0, 0.0, 0.0)
hit_geom_id = int(-1)
if local_id < bvh_ngeom:
gi = enabled_geom_ids[local_id]
gtype = geom_type[gi]
else:
gi = local_id - bvh_ngeom
gtype = GeomType.FLEX
hit_geom_id = gi
# TODO: Investigate branch elimination with static loop unrolling
if geom_type[gi] == GeomType.PLANE:
if gtype == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.HFIELD:
if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
)
if geom_type[gi] == GeomType.SPHERE:
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
geom_xpos_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.ELLIPSOID:
if gtype == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CAPSULE:
if gtype == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CYLINDER:
if gtype == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.BOX:
if gtype == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.MESH:
if gtype == GeomType.MESH:
d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
mesh_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
)
if gtype == GeomType.FLEX:
hit_geom_id = -2
flexid = flex_geom_flexid[gi]
edge_id = flex_geom_edgeid[gi]
if edge_id >= 0:
edge = flex_edge[edge_id]
vert_adr = flex_vertadr[flexid]
v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
pos = 0.5 * (v0 + v1)
vec = v1 - v0
length = wp.length(vec)
edgeq = math.quat_z2vec(vec)
mat = math.quat_to_mat(edgeq)
size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
hit_mesh_id = flexid
else:
flex_gr = flex_group_root[worldid, flexid]
d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist)
if d >= 0.0:
hit_mesh_id = flexid
if d >= 0.0 and d < dist:
dist = d
normal = n
geom_id = gi
geom_id = hit_geom_id
bary_u = u
bary_v = v
face_idx = f
@@ -217,17 +259,26 @@ def cast_ray_first_hit(
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
flex_vertadr: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_radius: wp.array(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
bvh_id: wp.uint64,
group_root: int,
world_id: int,
worldid: int,
bvh_ngeom: int,
bvh_nflexgeom: int,
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
flex_geom_flexid: wp.array(dtype=int),
flex_geom_edgeid: wp.array(dtype=int),
flex_bvh_id: wp.array(dtype=wp.uint64),
flex_group_root: wp.array2d(dtype=int),
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
max_dist: float,
@@ -235,81 +286,119 @@ def cast_ray_first_hit(
"""A simpler version of casting rays that only checks for the first hit."""
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
bounds_nr = int(0)
ngeom = bvh_ngeom + bvh_nflexgeom
while wp.bvh_query_next(query, bounds_nr, max_dist):
gi_global = bounds_nr
gi_bvh_local = gi_global - (world_id * bvh_ngeom)
gi = enabled_geom_ids[gi_bvh_local]
local_id = gi_global - (worldid * ngeom)
d = float(-1.0)
n = wp.vec3(0.0, 0.0, 0.0)
if local_id < bvh_ngeom:
gi = enabled_geom_ids[local_id]
gtype = geom_type[gi]
else:
gi = local_id - bvh_ngeom
gtype = GeomType.FLEX
# TODO: Investigate branch elimination with static loop unrolling
if geom_type[gi] == GeomType.PLANE:
if gtype == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.HFIELD:
if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
max_dist,
)
if geom_type[gi] == GeomType.SPHERE:
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
geom_xpos_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.ELLIPSOID:
if gtype == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CAPSULE:
if gtype == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CYLINDER:
if gtype == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.BOX:
if gtype == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.MESH:
if gtype == GeomType.MESH:
hit = ray_mesh_with_bvh_anyhit(
mesh_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
max_dist,
)
d = 0.0 if hit else -1.0
if gtype == GeomType.FLEX:
flexid = flex_geom_flexid[gi]
edge_id = flex_geom_edgeid[gi]
if edge_id >= 0:
edge = flex_edge[edge_id]
vert_adr = flex_vertadr[flexid]
v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
pos = 0.5 * (v0 + v1)
vec = v1 - v0
length = wp.length(vec)
edgeq = math.quat_z2vec(vec)
mat = math.quat_to_mat(edgeq)
size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
else:
hit = ray_flex_with_bvh_anyhit(
flex_bvh_id,
flexid,
flex_group_root[worldid, flexid],
ray_origin_world,
ray_dir_world,
max_dist,
)
d = 0.0 if hit else -1.0
if d >= 0.0 and d < max_dist:
return True
@@ -323,18 +412,27 @@ def compute_lighting(
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
flex_vertadr: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_radius: wp.array(dtype=float),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
use_shadows: bool,
bvh_id: wp.uint64,
group_root: int,
bvh_ngeom: int,
bvh_nflexgeom: int,
enabled_geom_ids: wp.array(dtype=int),
world_id: int,
worldid: int,
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
flex_geom_flexid: wp.array(dtype=int),
flex_geom_edgeid: wp.array(dtype=int),
flex_bvh_id: wp.array(dtype=wp.uint64),
flex_group_root: wp.array2d(dtype=int),
lightactive: bool,
lighttype: int,
lightcastshadow: bool,
@@ -385,15 +483,24 @@ def compute_lighting(
geom_type,
geom_dataid,
geom_size,
flex_vertadr,
flex_edge,
flex_radius,
geom_xpos_in,
geom_xmat_in,
flexvert_xpos_in,
bvh_id,
group_root,
world_id,
worldid,
bvh_ngeom,
bvh_nflexgeom,
enabled_geom_ids,
mesh_bvh_id,
hfield_bvh_id,
flex_geom_flexid,
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
shadow_origin,
L,
max_t,
@@ -418,6 +525,7 @@ def render(m: Model, d: Data, rc: RenderContext):
"""
rc.rgb_data.fill_(rc.background_color)
rc.depth_data.fill_(0.0)
rc.seg_data.fill_(-1)
@wp.kernel(module="unique", enable_backward=False)
def _render_megakernel(
@@ -434,6 +542,9 @@ def render(m: Model, d: Data, rc: RenderContext):
light_type: wp.array2d(dtype=int),
light_castshadow: wp.array2d(dtype=bool),
light_active: wp.array2d(dtype=bool),
flex_vertadr: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_radius: wp.array(dtype=float),
mesh_faceadr: wp.array(dtype=int),
mat_texid: wp.array3d(dtype=int),
mat_texrepeat: wp.array2d(dtype=wp.vec2),
@@ -445,21 +556,25 @@ def render(m: Model, d: Data, rc: RenderContext):
cam_xmat_in: wp.array2d(dtype=wp.mat33),
light_xpos_in: wp.array2d(dtype=wp.vec3),
light_xdir_in: wp.array2d(dtype=wp.vec3),
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# In:
nrender: int,
use_shadows: bool,
bvh_ngeom: int,
bvh_nflexgeom: int,
cam_res: wp.array(dtype=wp.vec2i),
cam_id_map: wp.array(dtype=int),
ray: wp.array(dtype=wp.vec3),
rgb_adr: wp.array(dtype=int),
depth_adr: wp.array(dtype=int),
seg_adr: wp.array(dtype=int),
render_rgb: wp.array(dtype=bool),
render_depth: wp.array(dtype=bool),
render_seg: wp.array(dtype=bool),
bvh_id: wp.uint64,
group_root: wp.array(dtype=int),
flex_bvh_id: wp.uint64,
flex_group_root: wp.array(dtype=int),
flex_bvh_id: wp.array(dtype=wp.uint64),
flex_group_root: wp.array2d(dtype=int),
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
mesh_facetexcoord: wp.array(dtype=wp.vec3i),
@@ -467,46 +582,48 @@ def render(m: Model, d: Data, rc: RenderContext):
mesh_texcoord_offsets: wp.array(dtype=int),
hfield_bvh_id: wp.array(dtype=wp.uint64),
flex_rgba: wp.array(dtype=wp.vec4),
# TODO: remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
textures: wp.array(dtype=TEXTURE_DTYPE),
flex_geom_flexid: wp.array(dtype=int),
flex_geom_edgeid: wp.array(dtype=int),
textures: wp.array(dtype=wp.Texture2D),
# Out:
rgb_out: wp.array2d(dtype=wp.uint32),
depth_out: wp.array2d(dtype=float),
seg_out: wp.array2d(dtype=int),
):
world_idx, ray_idx = wp.tid()
worldid, rayid = wp.tid()
# Map global ray_idx -> (cam_idx, ray_idx_local) using cumulative sizes
# Map global rayid -> (cam_idx, rayid_local) using cumulative sizes
cam_idx = int(-1)
ray_idx_local = int(-1)
rayid_local = int(-1)
accum = int(0)
for i in range(nrender):
num_i = cam_res[i][0] * cam_res[i][1]
if ray_idx < accum + num_i:
if rayid < accum + num_i:
cam_idx = i
ray_idx_local = ray_idx - accum
rayid_local = rayid - accum
break
accum += num_i
if cam_idx == -1 or ray_idx_local < 0:
if cam_idx == -1 or rayid_local < 0:
return
if not render_rgb[cam_idx] and not render_depth[cam_idx]:
if not render_rgb[cam_idx] and not render_depth[cam_idx] and not render_seg[cam_idx]:
return
# Map active camera index to MuJoCo camera ID
mujoco_cam_id = cam_id_map[cam_idx]
if wp.static(rc.use_precomputed_rays):
ray_dir_local_cam = ray[ray_idx]
ray_dir_local_cam = ray[rayid]
else:
img_w = cam_res[cam_idx][0]
img_h = cam_res[cam_idx][1]
px = ray_idx_local % img_w
py = ray_idx_local // img_w
px = rayid_local % img_w
py = rayid_local // img_w
ray_dir_local_cam = compute_ray(
cam_projection[mujoco_cam_id],
cam_fovy[world_idx % cam_fovy.shape[0], mujoco_cam_id],
cam_fovy[worldid % cam_fovy.shape[0], mujoco_cam_id],
cam_sensorsize[mujoco_cam_id],
cam_intrinsic[world_idx % cam_intrinsic.shape[0], mujoco_cam_id],
cam_intrinsic[worldid % cam_intrinsic.shape[0], mujoco_cam_id],
img_w,
img_h,
px,
@@ -514,38 +631,37 @@ def render(m: Model, d: Data, rc: RenderContext):
wp.static(rc.znear),
)
ray_dir_world = cam_xmat_in[world_idx, mujoco_cam_id] @ ray_dir_local_cam
ray_origin_world = cam_xpos_in[world_idx, mujoco_cam_id]
ray_dir_world = cam_xmat_in[worldid, mujoco_cam_id] @ ray_dir_local_cam
ray_origin_world = cam_xpos_in[worldid, mujoco_cam_id]
geom_id, dist, normal, u, v, f, mesh_id = cast_ray(
geom_type,
geom_dataid,
geom_size,
flex_vertadr,
flex_edge,
flex_radius,
geom_xpos_in,
geom_xmat_in,
flexvert_xpos_in,
bvh_id,
group_root[world_idx],
world_idx,
group_root[worldid],
worldid,
bvh_ngeom,
bvh_nflexgeom,
enabled_geom_ids,
mesh_bvh_id,
hfield_bvh_id,
flex_geom_flexid,
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
ray_origin_world,
ray_dir_world,
)
if wp.static(m.nflex > 0):
d, n, u, v, f = ray_flex_with_bvh(
flex_bvh_id,
flex_group_root[world_idx],
ray_origin_world,
ray_dir_world,
dist,
)
if d >= 0.0 and d < dist:
dist = d
normal = n
geom_id = -2
if render_seg[cam_idx] and geom_id != -1:
seg_out[worldid, seg_adr[cam_idx] + rayid_local] = geom_id
# Early Out
if geom_id == -1:
@@ -556,9 +672,7 @@ def render(m: Model, d: Data, rc: RenderContext):
# In camera-local coordinates, the optical axis is -Z. The Z-component of the
# normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ)
# between the ray and the optical axis.
depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist * (
-ray_dir_local_cam[2]
)
depth_out[worldid, depth_adr[cam_idx] + rayid_local] = dist * (-ray_dir_local_cam[2])
if not render_rgb[cam_idx]:
return
@@ -567,31 +681,30 @@ def render(m: Model, d: Data, rc: RenderContext):
hit_point = ray_origin_world + ray_dir_world * dist
if geom_id == -2:
# TODO: Currently flex textures are not supported, and only the first rgba value
# is used until further flex support is added.
color = flex_rgba[0]
elif geom_matid[world_idx % geom_matid.shape[0], geom_id] == -1:
color = geom_rgba[world_idx % geom_rgba.shape[0], geom_id]
# We encode flex_id in mesh_id for flex ray hits during cast_ray
color = flex_rgba[mesh_id]
elif geom_matid[worldid % geom_matid.shape[0], geom_id] == -1:
color = geom_rgba[worldid % geom_rgba.shape[0], geom_id]
else:
color = mat_rgba[world_idx % mat_rgba.shape[0], geom_matid[world_idx % geom_matid.shape[0], geom_id]]
color = mat_rgba[worldid % mat_rgba.shape[0], geom_matid[worldid % geom_matid.shape[0], geom_id]]
base_color = wp.vec3(color[0], color[1], color[2])
hit_color = base_color
if wp.static(rc.use_textures):
if geom_id != -2:
mat_id = geom_matid[world_idx % geom_matid.shape[0], geom_id]
mat_id = geom_matid[worldid % geom_matid.shape[0], geom_id]
if mat_id >= 0:
tex_id = mat_texid[world_idx % mat_texid.shape[0], mat_id, 1]
tex_id = mat_texid[worldid % mat_texid.shape[0], mat_id, 1]
if tex_id >= 0:
tex_color = sample_texture(
geom_type,
mesh_faceadr,
geom_id,
mat_texrepeat[world_idx % mat_texrepeat.shape[0], mat_id],
mat_texrepeat[worldid % mat_texrepeat.shape[0], mat_id],
textures[tex_id],
geom_xpos_in[world_idx, geom_id],
geom_xmat_in[world_idx, geom_id],
geom_xpos_in[worldid, geom_id],
geom_xmat_in[worldid, geom_id],
mesh_facetexcoord,
mesh_texcoord,
mesh_texcoord_offsets,
@@ -616,21 +729,30 @@ def render(m: Model, d: Data, rc: RenderContext):
geom_type,
geom_dataid,
geom_size,
flex_vertadr,
flex_edge,
flex_radius,
geom_xpos_in,
geom_xmat_in,
flexvert_xpos_in,
use_shadows,
bvh_id,
group_root[world_idx],
group_root[worldid],
bvh_ngeom,
bvh_nflexgeom,
enabled_geom_ids,
world_idx,
worldid,
mesh_bvh_id,
hfield_bvh_id,
light_active[world_idx % light_active.shape[0], l],
light_type[world_idx % light_type.shape[0], l],
light_castshadow[world_idx % light_castshadow.shape[0], l],
light_xpos_in[world_idx, l],
light_xdir_in[world_idx, l],
flex_geom_flexid,
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
light_active[worldid % light_active.shape[0], l],
light_type[worldid % light_type.shape[0], l],
light_castshadow[worldid % light_castshadow.shape[0], l],
light_xpos_in[worldid, l],
light_xdir_in[worldid, l],
normal,
hit_point,
)
@@ -639,7 +761,7 @@ def render(m: Model, d: Data, rc: RenderContext):
hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0))
hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0))
rgb_out[world_idx, rgb_adr[cam_idx] + ray_idx_local] = pack_rgba_to_uint32(
rgb_out[worldid, rgb_adr[cam_idx] + rayid_local] = pack_rgba_to_uint32(
hit_color[0] * 255.0,
hit_color[1] * 255.0,
hit_color[2] * 255.0,
@@ -662,6 +784,9 @@ def render(m: Model, d: Data, rc: RenderContext):
m.light_type,
m.light_castshadow,
m.light_active,
m.flex_vertadr,
m.flex_edge,
m.flex_radius,
m.mesh_faceadr,
m.mat_texid,
m.mat_texrepeat,
@@ -672,16 +797,20 @@ def render(m: Model, d: Data, rc: RenderContext):
d.cam_xmat,
d.light_xpos,
d.light_xdir,
d.flexvert_xpos,
rc.nrender,
rc.use_shadows,
rc.bvh_ngeom,
rc.bvh_nflexgeom,
rc.cam_res,
rc.cam_id_map,
rc.ray,
rc.rgb_adr,
rc.depth_adr,
rc.seg_adr,
rc.render_rgb,
rc.render_depth,
rc.render_seg,
rc.bvh_id,
rc.group_root,
rc.flex_bvh_id,
@@ -693,10 +822,13 @@ def render(m: Model, d: Data, rc: RenderContext):
rc.mesh_texcoord_offsets,
rc.hfield_bvh_id,
rc.flex_rgba,
rc.flex_geom_flexid,
rc.flex_geom_edgeid,
rc.textures,
],
outputs=[
rc.rgb_data,
rc.depth_data,
rc.seg_data,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
Vendored
+38
View File
@@ -206,3 +206,41 @@ def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_ou
inputs=[rc.depth_data, rc.depth_adr, camera_index, depth_scale],
outputs=[depth_out],
)
@wp.kernel
def _extract_seg_kernel(
# In:
seg_data: wp.array2d(dtype=int),
seg_adr: wp.array(dtype=int),
camera_index: int,
# Out:
seg_out: wp.array3d(dtype=int),
):
"""Extract per-pixel geom IDs from the render context buffers for a given camera index."""
worldid, pixelid = wp.tid()
xid = pixelid % seg_out.shape[2]
yid = pixelid // seg_out.shape[2]
seg_adr_offset = seg_adr[camera_index]
seg_out[worldid, yid, xid] = seg_data[worldid, seg_adr_offset + pixelid]
def get_segmentation(rc: RenderContext, camera_index: int, seg_out: wp.array3d(dtype=int)):
"""Get the segmentation data from the render context buffers for a given camera index.
Each pixel contains the MuJoCo geom ID of the geometry hit by the ray, -1 for
background, or -2 for flex bodies.
Args:
rc: The render context on device.
camera_index: The index of the camera to get the segmentation data for.
seg_out: The output array to store the geom IDs in, with shape
(nworld, height, width).
"""
wp.launch(
_extract_seg_kernel,
dim=(seg_out.shape[0], seg_out.shape[1] * seg_out.shape[2]),
inputs=[rc.seg_data, rc.seg_adr, camera_index],
outputs=[seg_out],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+123 -149
View File
@@ -15,12 +15,17 @@
from typing import Any, Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src import ray
from mujoco.mjx.third_party.mujoco_warp._src import smooth
from mujoco.mjx.third_party.mujoco_warp._src import support
from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import get_sdf_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
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 MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -28,9 +33,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DataType
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
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 MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType
@@ -40,10 +42,10 @@ 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 vec8
from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
from mujoco.mjx.third_party.mujoco_warp._src.types import vec_pluginattr
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
import warp as wp
wp.set_module_options({"enable_backward": False})
@@ -2081,16 +2083,16 @@ def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3:
@wp.kernel
def _preprocess_tactile_contacts(
# Model:
body_weldid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
# Data in:
contact_geom_in: wp.array(dtype=wp.vec2i),
contact_worldid_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
# Out:
weld_geom_count_out: wp.array2d(dtype=int),
weld_geom_list_out: wp.array3d(dtype=int),
# Model:
body_weldid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
# Data in:
contact_geom_in: wp.array(dtype=wp.vec2i),
contact_worldid_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
# Out:
weld_geom_count_out: wp.array2d(dtype=int),
weld_geom_list_out: wp.array3d(dtype=int),
):
conid = wp.tid()
ncon = nacon_in[0]
@@ -2118,42 +2120,43 @@ def _preprocess_tactile_contacts(
@wp.kernel
def _sensor_tactile(
# Model:
body_rootid: wp.array(dtype=int),
body_weldid: wp.array(dtype=int),
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
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),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_octadr: wp.array(dtype=int),
mesh_normaladr: wp.array(dtype=int),
mesh_normalnum: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_normal: wp.array(dtype=wp.vec3),
mesh_quat: wp.array(dtype=wp.quat),
sensor_objid: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_dim: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
geom_plugin_index: wp.array(dtype=int),
taxel_vertadr: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
weld_geom_count_in: wp.array2d(dtype=int),
weld_geom_list_in: wp.array3d(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
# Model:
body_rootid: wp.array(dtype=int),
body_weldid: wp.array(dtype=int),
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8),
geom_type: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_octadr: wp.array(dtype=int),
mesh_normaladr: wp.array(dtype=int),
mesh_normalnum: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_normal: wp.array(dtype=wp.vec3),
mesh_quat: wp.array(dtype=wp.quat),
sensor_objid: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_dim: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=vec_pluginattr),
geom_plugin_index: wp.array(dtype=int),
taxel_vertadr: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
weld_geom_count_in: wp.array2d(dtype=int),
weld_geom_list_in: wp.array3d(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, taxelid = wp.tid()
@@ -2211,40 +2214,25 @@ def _sensor_tactile(
contact_type = geom_type[geom]
plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params(
oct_child,
oct_aabb,
oct_coeff,
mesh_octadr,
plugin,
plugin_attr,
contact_type,
geom_size[worldid % geom_size.shape[0], geom],
plugin_id,
mesh_id,
oct_child,
oct_aabb,
oct_coeff,
mesh_octadr,
plugin,
plugin_attr,
contact_type,
geom_size[worldid % geom_size.shape[0], geom],
plugin_id,
geom_dataid[geom],
)
depth = wp.min(
sdf(
contact_type,
lpos,
plugin_attributes,
plugin_index,
volume_data,
mesh_data,
),
0.0,
)
depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0)
if depth >= 0.0:
continue
vel_sensor = _transform_spatial(
cvel_in[worldid, parent_weld],
xpos - subtree_com_in[worldid, body_rootid[parent_weld]],
)
vel_sensor = _transform_spatial(cvel_in[worldid, parent_weld], xpos - subtree_com_in[worldid, body_rootid[parent_weld]])
vel_other = _transform_spatial(
cvel_in[worldid, body],
geom_xpos_in[worldid, geom]
- subtree_com_in[worldid, body_rootid[body]],
cvel_in[worldid, body], geom_xpos_in[worldid, geom] - subtree_com_in[worldid, body_rootid[body]]
)
vel_rel = vel_sensor - vel_other
@@ -2259,24 +2247,9 @@ def _sensor_tactile(
forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
dim = sensor_dim[sensor_id] // 3
wp.atomic_add(
sensordata_out,
worldid,
sensor_adr[sensor_id] + 0 * dim + vertid,
forceT[0],
)
wp.atomic_add(
sensordata_out,
worldid,
sensor_adr[sensor_id] + 1 * dim + vertid,
forceT[1],
)
wp.atomic_add(
sensordata_out,
worldid,
sensor_adr[sensor_id] + 2 * dim + vertid,
forceT[2],
)
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
@wp.func
@@ -2507,60 +2480,61 @@ def sensor_acc(m: Model, d: Data):
weld_geom_count = wp.zeros((d.nworld, m.nbody), dtype=int)
weld_geom_list = wp.full((d.nworld, m.nbody, MJ_MAXCONPAIR), -1, dtype=int)
wp.launch(
_preprocess_tactile_contacts,
dim=d.naconmax,
inputs=[
m.body_weldid,
m.geom_bodyid,
d.contact.geom,
d.contact.worldid,
d.nacon,
],
outputs=[
weld_geom_count,
weld_geom_list,
],
_preprocess_tactile_contacts,
dim=d.naconmax,
inputs=[
m.body_weldid,
m.geom_bodyid,
d.contact.geom,
d.contact.worldid,
d.nacon,
],
outputs=[
weld_geom_count,
weld_geom_list,
],
)
wp.launch(
_sensor_tactile,
dim=(d.nworld, m.nsensortaxel),
inputs=[
m.body_rootid,
m.body_weldid,
m.oct_child,
m.oct_aabb,
m.oct_coeff,
m.geom_type,
m.geom_bodyid,
m.geom_size,
m.mesh_vertadr,
m.mesh_vertnum,
m.mesh_octadr,
m.mesh_normaladr,
m.mesh_normalnum,
m.mesh_vert,
m.mesh_normal,
m.mesh_quat,
m.sensor_objid,
m.sensor_refid,
m.sensor_dim,
m.sensor_adr,
m.plugin,
m.plugin_attr,
m.geom_plugin_index,
m.taxel_vertadr,
m.taxel_sensorid,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
weld_geom_count,
weld_geom_list,
],
outputs=[
d.sensordata,
],
_sensor_tactile,
dim=(d.nworld, m.nsensortaxel),
inputs=[
m.body_rootid,
m.body_weldid,
m.oct_child,
m.oct_aabb,
m.oct_coeff,
m.geom_type,
m.geom_bodyid,
m.geom_dataid,
m.geom_size,
m.mesh_vertadr,
m.mesh_vertnum,
m.mesh_octadr,
m.mesh_normaladr,
m.mesh_normalnum,
m.mesh_vert,
m.mesh_normal,
m.mesh_quat,
m.sensor_objid,
m.sensor_refid,
m.sensor_dim,
m.sensor_adr,
m.plugin,
m.plugin_attr,
m.geom_plugin_index,
m.taxel_vertadr,
m.taxel_sensorid,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
weld_geom_count,
weld_geom_list,
],
outputs=[
d.sensordata,
],
)
sensor_contact_nmatch = wp.empty((d.nworld, m.nsensorcontact), dtype=int)
@@ -2882,12 +2856,12 @@ def energy_pos(m: Model, d: Data):
wp.launch(_energy_pos_zero, dim=d.nworld, outputs=[d.energy])
# init potential energy: -sum_i(body_i.mass * dot(gravity, body_i.pos))
if not m.opt.disableflags & DisableBit.GRAVITY:
if not (m.opt.disableflags & DisableBit.GRAVITY):
wp.launch(
_energy_pos_gravity, dim=(d.nworld, m.nbody - 1), inputs=[m.opt.gravity, m.body_mass, d.xipos], outputs=[d.energy]
)
if not m.opt.disableflags & DisableBit.SPRING:
if not (m.opt.disableflags & DisableBit.SPRING):
# add joint-level springs
wp.launch(
_energy_pos_passive_joint,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+800 -552
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+514 -666
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+34
View File
@@ -18,18 +18,52 @@ from typing import Optional, Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.math import motion_cross
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import State
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
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
wp.set_module_options({"enable_backward": False})
# TODO(team): kernel analyzer array slice?
@wp.func
def next_act(
# Model:
opt_timestep: float, # kernel_analyzer: ignore
actuator_dyntype: int, # kernel_analyzer: ignore
actuator_dynprm: vec10f, # kernel_analyzer: ignore
actuator_actrange: wp.vec2, # kernel_analyzer: ignore
# Data In:
act_in: float, # kernel_analyzer: ignore
act_dot_in: float, # kernel_analyzer: ignore
# In:
act_dot_scale: float,
clamp: bool,
) -> float:
# advance actuation
if actuator_dyntype == DynType.FILTEREXACT:
tau = wp.max(MJ_MINVAL, actuator_dynprm[0])
act = act_in + act_dot_scale * act_dot_in * tau * (1.0 - wp.exp(-opt_timestep / tau))
elif actuator_dyntype == DynType.USER:
return act_in
else:
act = act_in + act_dot_scale * act_dot_in * opt_timestep
# clamp to actrange
if clamp:
act = wp.clamp(act, actuator_actrange[0], actuator_actrange[1])
return act
@cache_kernel
def mul_m_sparse(check_skip: bool):
@wp.kernel(module="unique")
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+137 -144
View File
@@ -33,11 +33,8 @@ MJ_MAX_EPAFACES = 5
TILE_SIZE_JTDAJ_SPARSE = 16
TILE_SIZE_JTDAJ_DENSE = 16
# TODO(team): remove after improving performance for sparse constraint jacobian
SPARSE_CONSTRAINT_JACOBIAN = False
# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
TEXTURE_DTYPE = wp.Texture2D if hasattr(wp, "Texture2D") else int
# maximum number of plugin attributes
_NPLUGINATTR = 128
# TODO(team): add check that all wp.launch_tiled 'block_dim' settings are configurable
@@ -53,7 +50,6 @@ class BlockDim:
# forward
euler_dense: int = 32
actuator_velocity: int = 32
tendon_velocity: int = 32
# ray
ray: int = 64
# sensor
@@ -63,6 +59,7 @@ class BlockDim:
cholesky_factorize: int = 32
cholesky_solve: int = 32
cholesky_factorize_solve: int = 32
solve_LD_sparse_fused: int = 64
# solver
update_gradient_cholesky: int = 64
update_gradient_cholesky_blocked: int = 32
@@ -351,6 +348,7 @@ class GeomType(enum.IntEnum):
BOX: box
MESH: mesh
SDF: sdf
FLEX: flex
"""
PLANE = mujoco.mjtGeom.mjGEOM_PLANE
@@ -362,6 +360,7 @@ class GeomType(enum.IntEnum):
BOX = mujoco.mjtGeom.mjGEOM_BOX
MESH = mujoco.mjtGeom.mjGEOM_MESH
SDF = mujoco.mjtGeom.mjGEOM_SDF
FLEX = mujoco.mjtGeom.mjGEOM_FLEX
# unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
@@ -662,6 +661,10 @@ class vec11f(wp.types.vector(length=11, dtype=float)):
pass
class vec_pluginattr(wp.types.vector(length=_NPLUGINATTR, dtype=float)):
pass
class mat23f(wp.types.matrix(shape=(2, 3), dtype=float)):
pass
@@ -679,6 +682,7 @@ vec6 = vec6f
vec8 = vec8f
vec10 = vec10f
vec11 = vec11f
vec128 = vec_pluginattr
mat23 = mat23f
mat43 = mat43f
mat63 = mat63f
@@ -841,6 +845,7 @@ class Model:
nflexelem: number of elements in all flexes
nflexelemdata: number of element vertex ids in all flexes
nflexelemedge: number of element edge ids in all flexes
nflexshelldata: number of shell fragment vertex ids in all flexes
nJfe: number of non-zeros in sparse flexedge Jacobian
nmesh: number of meshes
nmeshvert: number of vertices for all meshes
@@ -857,6 +862,7 @@ class Model:
nexclude: number of excluded geom pairs
neq: number of equality constraints
ntendon: number of tendons
nJten: number of non-zeros in sparse tendon Jacobian
nwrap: number of wrap objects in all tendon paths
nsensor: number of sensors
nmocap: number of mocap bodies
@@ -973,6 +979,11 @@ class Model:
light_poscom0: global position rel. to sub-com in qpos0 (*, nlight, 3)
light_pos0: global position rel. to body in qpos0 (*, nlight, 3)
light_dir0: global direction in qpos0 (*, nlight, 3)
flex_contype: flex contact type (nflex,)
flex_conaffinity: flex contact affinity (nflex,)
flex_condim: contact dimensionality (1, 3, 4, 6) (nflex,)
flex_friction: friction for (slide, spin, roll) (nflex, 3)
flex_margin: detect contact if dist<margin (nflex,)
flex_dim: 1: lines, 2: triangles, 3: tetrahedra (nflex,)
flex_vertadr: first vertex address (nflex,)
flex_vertnum: number of vertices (nflex,)
@@ -980,17 +991,24 @@ class Model:
flex_edgenum: number of edges (nflex,)
flex_elemadr: first element address (nflex,)
flex_elemnum: number of elements (nflex,)
flex_elemedgeadr: first element address (nflex,)
flex_elemdataadr: first element vertex id address (nflex,)
flex_elemedgeadr: first element edge id address (nflex,)
flex_shellnum: number of shells (nflex,)
flex_shelldataadr: first shell data address (nflex,)
flex_vertbodyid: vertex body ids (nflexvert,)
flex_edge: edge vertex ids (2 per edge) (nflexedge, 2)
flex_edgeflap: adjacent vertex ids (dim=2 only) (nflexedge, 2)
flex_elem: element vertex ids (dim+1 per elem) (nflexelemdata,)
flex_elemedge: element edge ids (nflexelemedge,)
flex_shell: shell fragment vertex ids (dim per frag) (nflexshelldata,)
flex_vert: vertex local positions (nflexvert, 3)
flexedge_length0: edge lengths in qpos0 (nflexedge,)
flexedge_invweight0: inv. inertia for the edge (nflexedge,)
flex_radius: radius around primitive element (nflex,)
flex_stiffness: finite element stiffness matrix (nflexelem, 21)
flex_bending: bending stiffness (nflexedge, 17)
flex_damping: Rayleigh's damping coefficient (nflex,)
flex_centered: flex vertices are centered at body origin (nflex,)
flexedge_J_rownnz: number of nonzeros in Jacobian row (nflexedge,)
flexedge_J_rowadr: row start address in colind array (nflexedge,)
flexedge_J_colind: column indices in sparse Jacobian (nJfe,)
@@ -1020,8 +1038,7 @@ class Model:
hfield_ncol: number of columns in grid (nhfield,)
hfield_adr: start address in hfield_data (nhfield,)
hfield_data: elevation data (nhfielddata,)
mat_texid: texture id for rendering (*, nmat,
mjNTEXROLE)
mat_texid: texture id for rendering (*, nmat, mjNTEXROLE)
mat_texrepeat: texture repeat for rendering (*, nmat, 2)
mat_rgba: rgba (*, nmat, 4)
pair_dim: contact dimensionality (npair,)
@@ -1044,16 +1061,15 @@ class Model:
eq_data: numeric data for constraint (*, neq, mjNEQDATA)
tendon_adr: address of first object in tendon's path (ntendon,)
tendon_num: number of objects in tendon's path (ntendon,)
ten_J_rownnz: number of non-zeros in each tendon row (ntendon,)
ten_J_rowadr: row start address for sparse ten_J (ntendon,)
ten_J_colind: column indices in sparse ten_J (nJten,)
tendon_limited: does tendon have length limits (ntendon,)
tendon_actfrclimited: does ten have actuator force limit (ntendon,)
tendon_solref_lim: constraint solver reference: limit (*, ntendon,
mjNREF)
tendon_solimp_lim: constraint solver impedance: limit (*, ntendon,
mjNIMP)
tendon_solref_fri: constraint solver reference: friction (*, ntendon,
mjNREF)
tendon_solimp_fri: constraint solver impedance: friction (*, ntendon,
mjNIMP)
tendon_solref_lim: constraint solver reference: limit (*, ntendon, mjNREF)
tendon_solimp_lim: constraint solver impedance: limit (*, ntendon, mjNIMP)
tendon_solref_fri: constraint solver reference: friction (*, ntendon, mjNREF)
tendon_solimp_fri: constraint solver impedance: friction (*, ntendon, mjNIMP)
tendon_range: tendon length limits (*, ntendon, 2)
tendon_actfrcrange: range of total actuator force (*, ntendon, 2)
tendon_margin: min distance for limit detection (*, ntendon)
@@ -1099,11 +1115,12 @@ class Model:
sensor_adr: address in sensor array (nsensor,)
sensor_cutoff: cutoff for real and positive; 0: ignore (nsensor,)
plugin: globally registered plugin slot number (nplugin,)
plugin_attr: config attributes of geom plugin (nplugin, 3)
plugin_attr: config attributes of geom plugin (nplugin, _NPLUGINATTR)
M_rownnz: number of non-zeros in each row of qM (nv,)
M_rowadr: index of each row in qM (nv,)
M_colind: column indices of non-zeros in qM (nM,)
mapM2M: index mapping from M (legacy) to M (CSR) (nC)
flex_vertflexid: flex id for each flex vertex (nflexvert,)
warp only fields:
callback: custom physics callbacks
@@ -1120,14 +1137,12 @@ class Model:
nmaxpolygon: maximum number of verts per polygon
nmaxmeshdeg: maximum number of polygons per vert
is_sparse: whether to use sparse representations
has_fluid: True if wind, density, or viscosity are non-zero at put_model
time
has_fluid: True if wind, density, or viscosity are non-zero at put_model time
has_sdf_geom: whether the model contains SDF geoms
block_dim: block dim options
body_tree: list of body ids by tree level
body_branches: flattened body ids for all branches
body_branch_start: start index in body_branches for each branch (nbranch +
1,)
body_branch_start: start index in body_branches for each branch (nbranch + 1,)
mocap_bodyid: id of body for mocap (nmocap,)
body_fluid_ellipsoid: does body use ellipsoid fluid (nbody,)
jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
@@ -1153,6 +1168,7 @@ class Model:
tendon_site_pair_adr: site pair tendon address
tendon_geom_adr: geom tendon address
tendon_limited_adr: addresses for limited tendons
max_ten_J_rownnz: maximum number of non-zeros in a tendon row
ten_wrapadr_site: wrap object starting address for sites
ten_wrapnum_site: number of site wrap objects per tendon
wrap_jnt_adr: addresses for joint tendon wrap object
@@ -1189,6 +1205,8 @@ class Model:
taxel_sensorid: address for tactile sensors
qM_tiles: tiling configuration
qLD_updates: tuple of index triples for sparse factorization
qLD_all_updates: tuple of all levels concatenated
qLD_level_offsets: tuple of start offsets for each level
qM_fullm_i: sparse mass matrix addressing
qM_fullm_j: sparse mass matrix addressing
qM_mulm_rowadr: sparse matmul row pointers
@@ -1216,6 +1234,7 @@ class Model:
nflexelem: int
nflexelemdata: int
nflexelemedge: int
nflexshelldata: int
nJfe: int
nmesh: int
nmeshvert: int
@@ -1232,6 +1251,7 @@ class Model:
nexclude: int
neq: int
ntendon: int
nJten: int
nwrap: int
nsensor: int
nmocap: int
@@ -1348,6 +1368,11 @@ class Model:
light_poscom0: array("*", "nlight", wp.vec3)
light_pos0: array("*", "nlight", wp.vec3)
light_dir0: array("*", "nlight", wp.vec3)
flex_contype: array("nflex", int)
flex_conaffinity: array("nflex", int)
flex_condim: array("nflex", int)
flex_friction: array("nflex", wp.vec3)
flex_margin: array("nflex", float)
flex_dim: array("nflex", int)
flex_vertadr: array("nflex", int)
flex_vertnum: array("nflex", int)
@@ -1355,17 +1380,24 @@ class Model:
flex_edgenum: array("nflex", int)
flex_elemadr: array("nflex", int)
flex_elemnum: array("nflex", int)
flex_elemdataadr: array("nflex", int)
flex_elemedgeadr: array("nflex", int)
flex_shellnum: array("nflex", int)
flex_shelldataadr: array("nflex", int)
flex_vertbodyid: array("nflexvert", int)
flex_edge: array("nflexedge", wp.vec2i)
flex_edgeflap: array("nflexedge", wp.vec2i)
flex_elem: array("nflexelemdata", int)
flex_elemedge: array("nflexelemedge", int)
flex_shell: array("nflexshelldata", int)
flex_vert: array("nflexvert", wp.vec3)
flexedge_length0: array("nflexedge", float)
flexedge_invweight0: array("nflexedge", float)
flex_radius: array("nflex", float)
flex_stiffness: array("nflexelem", 21, float)
flex_bending: array("nflexedge", 17, float)
flex_damping: array("nflex", float)
flex_centered: array("nflex", bool)
flexedge_J_rownnz: array("nflexedge", int)
flexedge_J_rowadr: array("nflexedge", int)
flexedge_J_colind: array("nJfe", int)
@@ -1418,6 +1450,9 @@ class Model:
eq_data: array("*", "neq", vec11)
tendon_adr: array("ntendon", int)
tendon_num: array("ntendon", int)
ten_J_rownnz: array("ntendon", int)
ten_J_rowadr: array("ntendon", int)
ten_J_colind: array("nJten", int)
tendon_limited: array("ntendon", int)
tendon_actfrclimited: array("ntendon", bool)
tendon_solref_lim: array("*", "ntendon", wp.vec2)
@@ -1469,11 +1504,12 @@ class Model:
sensor_adr: array("nsensor", int)
sensor_cutoff: array("nsensor", float)
plugin: array("nplugin", int)
plugin_attr: array("nplugin", wp.vec3f)
plugin_attr: array("nplugin", vec_pluginattr)
M_rownnz: array("nv", int)
M_rowadr: array("nv", int)
M_colind: array("nC", int)
mapM2M: array("nC", int)
flex_vertflexid: array("nflexvert", int)
# warp only fields:
callback: Callback
nbranch: int
@@ -1515,6 +1551,7 @@ class Model:
tendon_site_pair_adr: wp.array(dtype=int)
tendon_geom_adr: wp.array(dtype=int)
tendon_limited_adr: wp.array(dtype=int)
max_ten_J_rownnz: int
ten_wrapadr_site: wp.array(dtype=int)
ten_wrapnum_site: wp.array(dtype=int)
wrap_jnt_adr: wp.array(dtype=int)
@@ -1546,6 +1583,8 @@ class Model:
taxel_sensorid: wp.array(dtype=int)
qM_tiles: tuple[TileSet, ...]
qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...]
qLD_all_updates: wp.array(dtype=wp.vec3i)
qLD_level_offsets: wp.array(dtype=int)
qM_fullm_i: wp.array(dtype=int)
qM_fullm_j: wp.array(dtype=int)
# Gather-based sparse mul_m indices (thread per DOF, no atomics)
@@ -1598,6 +1637,8 @@ class Contact:
solimp: array("naconmax", vec5)
dim: array("naconmax", int)
geom: array("naconmax", wp.vec2i)
flex: array("naconmax", wp.vec2i)
vert: array("naconmax", wp.vec2i)
efc_address: array("naconmax", "nmaxpyramid", int)
worldid: array("naconmax", int)
type: array("naconmax", int)
@@ -1611,14 +1652,14 @@ class Constraint:
Attributes:
type: constraint type (ConstraintType) (nworld, njmax)
id: id of object of specific type (nworld, njmax)
J_rownnz: number of non-zeros in J row (nworld, 0) dense (nworld,
njmax) sparse
J_rowadr: row start address in colind array (nworld, 0) dense (nworld,
njmax) sparse
J_rownnz: number of non-zeros in J row (nworld, 0) dense
(nworld, njmax) sparse
J_rowadr: row start address in colind array (nworld, 0) dense
(nworld, njmax) sparse
J_colind: column indices in J (nworld, 0, 0) dense
(nworld, 1, njmax * nv) sparse
J: constraint Jacobian (nworld, njmax_pad,
nv_pad) dense (nworld, 1, njmax * nv) sparse
(nworld, 1, njmax * nv) sparse
J: constraint Jacobian (nworld, njmax_pad, nv_pad) dense
(nworld, 1, njmax * nv) sparse
pos: constraint position (equality, contact) (nworld, njmax)
margin: inclusion margin (contact) (nworld, njmax)
D: constraint mass (nworld, njmax_pad)
@@ -1627,7 +1668,6 @@ class Constraint:
frictionloss: frictionloss (friction) (nworld, njmax)
force: constraint force in constraint space (nworld, njmax)
state: constraint state (nworld, njmax_pad)
warp only fields:
Ma: M*qacc (nworld, nv)
"""
@@ -1668,88 +1708,55 @@ class Data:
qacc_warmstart: acceleration used for warmstart (nworld, nv)
ctrl: control (nworld, nu)
qfrc_applied: applied generalized force (nworld, nv)
xfrc_applied: applied Cartesian force/torque (nworld, nbody,
6)
xfrc_applied: applied Cartesian force/torque (nworld, nbody, 6)
eq_active: enable/disable constraints (nworld, neq)
mocap_pos: position of mocap bodies (nworld, nmocap,
3)
mocap_quat: orientation of mocap bodies (nworld, nmocap,
4)
mocap_pos: position of mocap bodies (nworld, nmocap, 3)
mocap_quat: orientation of mocap bodies (nworld, nmocap, 4)
qacc: acceleration (nworld, nv)
act_dot: time-derivative of actuator activation (nworld, na)
sensordata: sensor data array (nworld,
nsensordata,)
xpos: Cartesian position of body frame (nworld, nbody,
3)
xquat: Cartesian orientation of body frame (nworld, nbody,
4)
xmat: Cartesian orientation of body frame (nworld, nbody,
3, 3)
xipos: Cartesian position of body com (nworld, nbody,
3)
ximat: Cartesian orientation of body inertia (nworld, nbody,
3, 3)
xanchor: Cartesian position of joint anchor (nworld, njnt,
3)
xaxis: Cartesian joint axis (nworld, njnt,
3)
geom_xpos: Cartesian geom position (nworld, ngeom,
3)
geom_xmat: Cartesian geom orientation (nworld, ngeom,
3, 3)
site_xpos: Cartesian site position (nworld, nsite,
3)
site_xmat: Cartesian site orientation (nworld, nsite,
3, 3)
cam_xpos: Cartesian camera position (nworld, ncam,
3)
cam_xmat: Cartesian camera orientation (nworld, ncam,
3, 3)
light_xpos: Cartesian light position (nworld, nlight,
3)
light_xdir: Cartesian light direction (nworld, nlight,
3)
subtree_com: center of mass of each subtree (nworld, nbody,
3)
sensordata: sensor data array (nworld, nsensordata,)
xpos: Cartesian position of body frame (nworld, nbody, 3)
xquat: Cartesian orientation of body frame (nworld, nbody, 4)
xmat: Cartesian orientation of body frame (nworld, nbody, 3, 3)
xipos: Cartesian position of body com (nworld, nbody, 3)
ximat: Cartesian orientation of body inertia (nworld, nbody, 3, 3)
xanchor: Cartesian position of joint anchor (nworld, njnt, 3)
xaxis: Cartesian joint axis (nworld, njnt, 3)
geom_xpos: Cartesian geom position (nworld, ngeom, 3)
geom_xmat: Cartesian geom orientation (nworld, ngeom, 3, 3)
site_xpos: Cartesian site position (nworld, nsite, 3)
site_xmat: Cartesian site orientation (nworld, nsite, 3, 3)
cam_xpos: Cartesian camera position (nworld, ncam, 3)
cam_xmat: Cartesian camera orientation (nworld, ncam, 3, 3)
light_xpos: Cartesian light position (nworld, nlight, 3)
light_xdir: Cartesian light direction (nworld, nlight, 3)
subtree_com: center of mass of each subtree (nworld, nbody, 3)
cdof: com-based motion axis of each dof (rot:lin) (nworld, nv, 6)
cinert: com-based body inertia and mass (nworld, nbody,
10)
flexvert_xpos: cartesian flex vertex positions (nworld,
nflexvert, 3)
cinert: com-based body inertia and mass (nworld, nbody, 10)
flexvert_xpos: cartesian flex vertex positions (nworld, nflexvert, 3)
flexedge_J: edge length Jacobian (nworld, nJfe)
flexedge_length: flex edge lengths (nworld,
nflexedge, 1)
ten_wrapadr: start address of tendon's path (nworld,
ntendon)
ten_wrapnum: number of wrap points in path (nworld,
ntendon)
ten_J: tendon Jacobian (nworld,
ntendon, nv)
ten_length: tendon lengths (nworld,
ntendon)
wrap_obj: geomid; -1: site; -2: pulley (nworld, nwrap,
2)
wrap_xpos: Cartesian 3D points in all paths (nworld, nwrap,
6)
flexedge_length: flex edge lengths (nworld, nflexedge, 1)
ten_wrapadr: start address of tendon's path (nworld, ntendon)
ten_wrapnum: number of wrap points in path (nworld, ntendon)
ten_J: tendon Jacobian (nworld, nJten)
ten_length: tendon lengths (nworld, ntendon)
wrap_obj: geomid; -1: site; -2: pulley (nworld, nwrap, 2)
wrap_xpos: Cartesian 3D points in all paths (nworld, nwrap, 6)
actuator_length: actuator lengths (nworld, nu)
moment_rownnz: number of non-zeros in actuator_moment row (nworld, nu)
moment_rowadr: row start address in actuator_moment (nworld, nu)
moment_colind: column indices in sparse actuator_moment (nworld, nJmom)
actuator_moment: actuator moments (nworld, nJmom)
crb: com-based composite inertia and mass (nworld, nbody,
10)
qM: total inertia (nworld, nv, nv)
if dense (nworld, 1, nM) if sparse
qLD: L'*D*L factorization of M (nworld, nv, nv)
if dense (nworld, 1, nC) if sparse
crb: com-based composite inertia and mass (nworld, nbody, 10)
qM: total inertia (nworld, nv, nv) if dense
(nworld, 1, nM) if sparse
qLD: L'*D*L factorization of M (nworld, nv, nv) if dense
(nworld, 1, nC) if sparse
qLDiagInv: 1/diag(D) (nworld, nv)
flexedge_velocity: flex edge velocities (nworld,
nflexedge)
ten_velocity: tendon velocities (nworld,
ntendon)
flexedge_velocity: flex edge velocities (nworld, nflexedge)
ten_velocity: tendon velocities (nworld, ntendon)
actuator_velocity: actuator velocities (nworld, nu)
cvel: com-based velocity (rot:lin) (nworld, nbody,
6)
cvel: com-based velocity (rot:lin) (nworld, nbody, 6)
cdof_dot: time-derivative of cdof (rot:lin) (nworld, nv, 6)
qfrc_bias: C(qpos,qvel) (nworld, nv)
qfrc_spring: passive spring force (nworld, nv)
@@ -1757,23 +1764,19 @@ class Data:
qfrc_gravcomp: passive gravity compensation force (nworld, nv)
qfrc_fluid: passive fluid force (nworld, nv)
qfrc_passive: total passive force (nworld, nv)
subtree_linvel: linear velocity of subtree com (nworld, nbody,
3)
subtree_angmom: angular momentum about subtree com (nworld, nbody,
3)
subtree_linvel: linear velocity of subtree com (nworld, nbody, 3)
subtree_angmom: angular momentum about subtree com (nworld, nbody, 3)
actuator_force: actuator force in actuation space (nworld, nu)
qfrc_actuator: actuator force (nworld, nv)
qfrc_smooth: net unconstrained force (nworld, nv)
qacc_smooth: unconstrained acceleration (nworld, nv)
qfrc_constraint: constraint force (nworld, nv)
qfrc_inverse: net external force; should equal: (nworld, nv)
qfrc_applied + J.T @ xfrc_applied + qfrc_actuator
cacc: com-based acceleration (nworld, nbody,
6)
cfrc_int: com-based interaction force with parent (nworld, nbody,
6)
cfrc_ext: com-based external force on body (nworld, nbody,
6)
qfrc_applied + J.T @ xfrc_applied
+ qfrc_actuator
cacc: com-based acceleration (nworld, nbody, 6)
cfrc_int: com-based interaction force with parent (nworld, nbody, 6)
cfrc_ext: com-based external force on body (nworld, nbody, 6)
contact: contact data
efc: constraint data
tree_island: island ID per tree (-1 if unconstrained) (nworld, ntree)
@@ -1784,6 +1787,7 @@ class Data:
naccdmax: maximum number of contacts for CCD (all worlds)
njmax: maximum number of constraints per world
njmax_pad: njmax rounded up to the nearest multiple of TILE_SIZE_JTDAJ
njmax_nnz: number of non-zeros in constraint Jacobian
nacon: number of detected contacts (across all worlds) (1,)
ncollision: collision count from broadphase (1,)
"""
@@ -1832,7 +1836,7 @@ class Data:
flexedge_length: array("nworld", "nflexedge", float)
ten_wrapadr: array("nworld", "ntendon", int)
ten_wrapnum: array("nworld", "ntendon", int)
ten_J: array("nworld", "ntendon", "nv", float)
ten_J: array("nworld", "nJten", float)
ten_length: array("nworld", "ntendon", float)
wrap_obj: array("nworld", "nwrap", wp.vec2i)
wrap_xpos: array("nworld", "nwrap", wp.spatial_vector)
@@ -1877,6 +1881,7 @@ class Data:
naccdmax: int
njmax: int
njmax_pad: int
njmax_nnz: int
nacon: array(1, int)
ncollision: array(1, int)
@@ -1905,21 +1910,12 @@ class RenderContext:
hfield_registry: hfield BVH id to warp mesh mapping
hfield_bvh_id: hfield BVH ids
hfield_bounds_size: hfield bounds half-extents
flex_mesh: flex mesh
flex_mesh_registry: per-flex mesh BVH registry (prevents garbage collection)
flex_rgba: flex rgba
flex_bvh_id: flex BVH id
flex_face_point: flex face points
flex_faceadr: flex face addresses
flex_nface: number of flex faces
flex_nwork: total flex work items for refit
flex_group_root: flex group roots
flex_elemdataadr: flex element data addresses
flex_shell: flex shell data
flex_shelldataadr: flex shell data addresses
flex_radius: flex radius
flex_workadr: flex work item addresses for refit
flex_worknum: flex work item counts for refit
flex_bvh_id: per-flex BVH ids
flex_group_root: per-flex group roots (nworld x n_flex_bvh)
flex_render_smooth: whether to render flex meshes smoothly
flex_dim: flex dimension per flex (1D/2D/3D)
bvh: scene BVH
bvh_id: scene BVH id
lower: lower bounds
@@ -1929,12 +1925,13 @@ class RenderContext:
ray: rays
rgb_data: RGB data
rgb_adr: RGB addresses
rgb_size: per-camera RGB buffer sizes
depth_data: depth data
depth_adr: depth addresses
depth_size: per-camera depth buffer sizes
render_rgb: per-camera RGB render flags
render_depth: per-camera depth render flags
seg_data: segmentation data (per-pixel geom IDs)
seg_adr: segmentation addresses
render_seg: per-camera segmentation render flags
znear: near plane distance
total_rays: total number of rays
"""
@@ -1954,27 +1951,20 @@ class RenderContext:
mesh_texcoord: array("*", wp.vec2)
mesh_texcoord_offsets: array("nmesh", int)
mesh_facetexcoord: array("nmeshface", wp.vec3i)
# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
textures: array("*", TEXTURE_DTYPE)
textures_registry: list[TEXTURE_DTYPE]
textures: array("*", wp.Texture2D)
textures_registry: list[wp.Texture2D]
hfield_registry: dict
hfield_bvh_id: array("nhfield", wp.uint64)
hfield_bounds_size: array("nhfield", wp.vec3)
flex_mesh: wp.Mesh
flex_mesh_registry: dict
flex_rgba: array("nflex", wp.vec4)
flex_bvh_id: wp.uint64
flex_face_point: array("*", wp.vec3)
flex_faceadr: array("nflex", int)
flex_nface: int
flex_nwork: int
flex_group_root: array("nworld", int)
flex_elemdataadr: array("nflex", int)
flex_shell: array("*", int)
flex_shelldataadr: array("nflex", int)
flex_radius: array("nflex", float)
flex_workadr: array("nflex", int)
flex_worknum: array("nflex", int)
flex_bvh_id: array("*", wp.uint64)
flex_group_root: array("nworld", "*", int)
flex_render_smooth: bool
bvh_nflexgeom: int
flex_dim_np: array("nflex", int)
flex_geom_flexid: array("*", int)
flex_geom_edgeid: array("*", int)
bvh: wp.Bvh
bvh_id: wp.uint64
lower: array("*", wp.vec3)
@@ -1988,5 +1978,8 @@ class RenderContext:
depth_adr: array("ncam", int)
render_rgb: array("ncam", bool)
render_depth: array("ncam", bool)
seg_data: array("*", int)
seg_adr: array("ncam", int)
render_seg: array("ncam", bool)
znear: float
total_rays: int
mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
Vendored
+8 -12
View File
@@ -37,9 +37,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
"""
# Split on both '.' and '-'
parts = re.split(r"[.\-]", version_str)
return tuple(
[(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)]
)
return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])
def check_version(spec: str) -> bool:
@@ -65,9 +63,7 @@ def check_version(spec: str) -> bool:
"""
match = re.match(r"^([a-zA-Z0-9_\-]+)(>=|<=|>|<|==|!=)(.+)$", spec)
if not match:
raise ValueError(
f"Invalid version spec '{spec}'. Expected format: 'package>=version'"
)
raise ValueError(f"Invalid version spec '{spec}'. Expected format: 'package>=version'")
package_name, op, version_str = match.groups()
required_version = _parse_version(version_str)
@@ -87,11 +83,11 @@ def check_version(spec: str) -> bool:
installed_version = _parse_version(installed_str)
ops = {
">=": operator.ge,
"<=": operator.le,
">": operator.gt,
"<": operator.lt,
"==": operator.eq,
"!=": operator.ne,
">=": operator.ge,
"<=": operator.le,
">": operator.gt,
"<": operator.lt,
"==": operator.eq,
"!=": operator.ne,
}
return ops[op](installed_version, required_version)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
Vendored
+2 -2
View File
@@ -146,13 +146,13 @@ def check_toolkit_driver():
if wp.get_device().is_cuda:
if not wp.is_conditional_graph_supported():
warnings.warn(
"""
"""
CUDA version < 12.4 detected
- graph capture may be unreliable for < 12.3
- conditional graph nodes are not available for < 12.4
Model.opt.graph_conditional should be set to False
""",
stacklevel=2,
stacklevel=2,
)
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+2 -2
View File
@@ -28,7 +28,7 @@ requires-python = ">=3.10"
dependencies = [
"absl-py",
"etils[epath]",
"mujoco>=3.5.0",
"mujoco>=3.6.0",
"numpy",
"warp-lang>=1.12",
]
@@ -55,7 +55,7 @@ dev = [
"ruff",
"pygls>=1.0.0,<2.0.0",
"lsprotocol>=2023.0.1,<2024.0.0",
"mujoco>=3.5.0.dev0",
"mujoco>=3.6.0.dev0",
"warp-lang>=1.11.0.dev0",
]
# TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+3 -2
View File
@@ -18,7 +18,7 @@
Usage: mjwarp-viewer <mjcf XML path> [flags]
Example:
mjwarp-viewer benchmark/humanoid/humanoid.xml -o "opt.solver=cg"
mjwarp-viewer benchmarks/humanoid/humanoid.xml -o "opt.solver=cg"
"""
import copy
@@ -56,6 +56,7 @@ _CLEAR_WARP_CACHE = flags.DEFINE_bool("clear_warp_cache", False, "Clear warp cac
_ENGINE = flags.DEFINE_enum_class("engine", EngineOptions.WARP, EngineOptions, "Simulation engine")
_NCONMAX = flags.DEFINE_integer("nconmax", None, "Maximum number of contacts.")
_NJMAX = flags.DEFINE_integer("njmax", None, "Maximum number of constraints per world.")
_NJMAX_NNZ = flags.DEFINE_integer("njmax_nnz", None, "Maximum number of non-zeros in constraint Jacobian.")
_NCCDMAX = flags.DEFINE_integer("nccdmax", None, "Maximum number of CCD contacts per world.")
_OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation: foo.bar = baz)", short_name="o")
_KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
@@ -149,7 +150,7 @@ def _main(argv: Sequence[str]) -> None:
override_model(mjm, _OVERRIDE.value)
m = mjw.put_model(mjm)
override_model(m, _OVERRIDE.value)
d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value, nccdmax=_NCCDMAX.value)
d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value, njmax_nnz=_NJMAX_NNZ.value, nccdmax=_NCCDMAX.value)
graph = _compile_step(m, d) if wp.get_device().is_cuda else None
if graph is None:
mjw.step(m, d) # warmup step
mjx/mujoco/mjx/warp/bvh.py
+25 -6
View File
@@ -48,20 +48,27 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _refit_bvh_shim(
# Model
nworld: int,
flex_dim: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_elem: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_radius: wp.array(dtype=float),
flex_shell: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_vertnum: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_type: wp.array(dtype=int),
nflex: int,
nflexelemdata: int,
nflexvert: int,
nflexelem: int,
# Data
flexvert_xpos: wp.array2d(dtype=wp.vec3),
geom_xmat: wp.array2d(dtype=wp.mat33),
@@ -77,15 +84,21 @@ def _refit_bvh_shim(
_d.efc = _e
_d.contact = _c
_m.flex_dim = flex_dim
_m.flex_edge = flex_edge
_m.flex_elem = flex_elem
_m.flex_elemadr = flex_elemadr
_m.flex_elemdataadr = flex_elemdataadr
_m.flex_elemnum = flex_elemnum
_m.flex_radius = flex_radius
_m.flex_shell = flex_shell
_m.flex_shelldataadr = flex_shelldataadr
_m.flex_vertadr = flex_vertadr
_m.flex_vertnum = flex_vertnum
_m.geom_dataid = geom_dataid
_m.geom_size = geom_size
_m.geom_type = geom_type
_m.nflex = nflex
_m.nflexelemdata = nflexelemdata
_m.nflexvert = nflexvert
_m.nflexelem = nflexelem
_d.flexvert_xpos = flexvert_xpos
_d.geom_xmat = geom_xmat
_d.geom_xpos = geom_xpos
@@ -113,15 +126,21 @@ def _refit_bvh_jax_impl(
out = jf(
d.qpos.shape[0],
m._impl.flex_dim,
m._impl.flex_edge,
m._impl.flex_elem,
m._impl.flex_elemadr,
m._impl.flex_elemdataadr,
m._impl.flex_elemnum,
m._impl.flex_radius,
m._impl.flex_shell,
m._impl.flex_shelldataadr,
m._impl.flex_vertadr,
m._impl.flex_vertnum,
m.geom_dataid,
m.geom_size,
m.geom_type,
m._impl.nflex,
m._impl.nflexelemdata,
m._impl.nflexvert,
m._impl.nflexelem,
d._impl.flexvert_xpos,
d.geom_xmat,
d.geom_xpos,
mjx/mujoco/mjx/warp/collision_driver.py
+100 -13
View File
@@ -52,8 +52,26 @@ def _collision_shim(
# Model
nworld: int,
block_dim: mjwp_types.BlockDim,
flex_conaffinity: wp.array(dtype=int),
flex_condim: wp.array(dtype=int),
flex_contype: wp.array(dtype=int),
flex_dim: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_friction: wp.array(dtype=wp.vec3),
flex_margin: wp.array(dtype=float),
flex_radius: wp.array(dtype=float),
flex_shell: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_shellnum: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_vertflexid: wp.array(dtype=int),
geom_aabb: wp.array3d(dtype=wp.vec3),
geom_conaffinity: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_contype: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_gap: wp.array2d(dtype=float),
@@ -90,6 +108,10 @@ def _collision_shim(
mesh_vert: wp.array(dtype=wp.vec3),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
nflex: int,
nflexelem: int,
nflexshelldata: int,
nflexvert: int,
ngeom: int,
nmaxmeshdeg: int,
nmaxpolygon: int,
@@ -108,7 +130,7 @@ def _collision_shim(
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
opt__broadphase: int,
opt__broadphase_filter: int,
opt__ccd_iterations: int,
@@ -120,12 +142,15 @@ def _collision_shim(
# Data
naccdmax: int,
naconmax: int,
flexvert_xpos: wp.array2d(dtype=wp.vec3),
geom_xmat: wp.array2d(dtype=wp.mat33),
geom_xpos: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int),
contact__dim: wp.array(dtype=int),
contact__dist: wp.array(dtype=float),
contact__efc_address: wp.array2d(dtype=int),
contact__flex: wp.array(dtype=wp.vec2i),
contact__frame: wp.array(dtype=wp.mat33),
contact__friction: wp.array(dtype=mjwp_types.vec5),
contact__geom: wp.array(dtype=wp.vec2i),
@@ -136,6 +161,7 @@ def _collision_shim(
contact__solref: wp.array(dtype=wp.vec2),
contact__solreffriction: wp.array(dtype=wp.vec2),
contact__type: wp.array(dtype=int),
contact__vert: wp.array(dtype=wp.vec2i),
contact__worldid: wp.array(dtype=int),
):
_m.stat = _s
@@ -144,8 +170,26 @@ def _collision_shim(
_d.efc = _e
_d.contact = _c
_m.block_dim = block_dim
_m.flex_conaffinity = flex_conaffinity
_m.flex_condim = flex_condim
_m.flex_contype = flex_contype
_m.flex_dim = flex_dim
_m.flex_elem = flex_elem
_m.flex_elemadr = flex_elemadr
_m.flex_elemdataadr = flex_elemdataadr
_m.flex_elemnum = flex_elemnum
_m.flex_friction = flex_friction
_m.flex_margin = flex_margin
_m.flex_radius = flex_radius
_m.flex_shell = flex_shell
_m.flex_shelldataadr = flex_shelldataadr
_m.flex_shellnum = flex_shellnum
_m.flex_vertadr = flex_vertadr
_m.flex_vertflexid = flex_vertflexid
_m.geom_aabb = geom_aabb
_m.geom_conaffinity = geom_conaffinity
_m.geom_condim = geom_condim
_m.geom_contype = geom_contype
_m.geom_dataid = geom_dataid
_m.geom_friction = geom_friction
_m.geom_gap = geom_gap
@@ -182,6 +226,10 @@ def _collision_shim(
_m.mesh_vert = mesh_vert
_m.mesh_vertadr = mesh_vertadr
_m.mesh_vertnum = mesh_vertnum
_m.nflex = nflex
_m.nflexelem = nflexelem
_m.nflexshelldata = nflexshelldata
_m.nflexvert = nflexvert
_m.ngeom = ngeom
_m.nmaxmeshdeg = nmaxmeshdeg
_m.nmaxpolygon = nmaxpolygon
@@ -211,6 +259,8 @@ def _collision_shim(
_m.plugin_attr = plugin_attr
_d.contact.dim = contact__dim
_d.contact.dist = contact__dist
_d.contact.efc_address = contact__efc_address
_d.contact.flex = contact__flex
_d.contact.frame = contact__frame
_d.contact.friction = contact__friction
_d.contact.geom = contact__geom
@@ -221,7 +271,9 @@ def _collision_shim(
_d.contact.solref = contact__solref
_d.contact.solreffriction = contact__solreffriction
_d.contact.type = contact__type
_d.contact.vert = contact__vert
_d.contact.worldid = contact__worldid
_d.flexvert_xpos = flexvert_xpos
_d.geom_xmat = geom_xmat
_d.geom_xpos = geom_xpos
_d.naccdmax = naccdmax
@@ -238,6 +290,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'ncollision': d._impl.ncollision.shape,
'contact__dim': d._impl.contact__dim.shape,
'contact__dist': d._impl.contact__dist.shape,
'contact__efc_address': d._impl.contact__efc_address.shape,
'contact__flex': d._impl.contact__flex.shape,
'contact__frame': d._impl.contact__frame.shape,
'contact__friction': d._impl.contact__friction.shape,
'contact__geom': d._impl.contact__geom.shape,
@@ -248,11 +302,12 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'contact__solref': d._impl.contact__solref.shape,
'contact__solreffriction': d._impl.contact__solreffriction.shape,
'contact__type': d._impl.contact__type.shape,
'contact__vert': d._impl.contact__vert.shape,
'contact__worldid': d._impl.contact__worldid.shape,
}
jf = ffi.jax_callable_variadic_tuple(
_collision_shim,
num_outputs=15,
num_outputs=18,
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([
@@ -260,6 +315,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'ncollision',
'contact__dim',
'contact__dist',
'contact__efc_address',
'contact__flex',
'contact__frame',
'contact__friction',
'contact__geom',
@@ -270,6 +327,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'contact__solref',
'contact__solreffriction',
'contact__type',
'contact__vert',
'contact__worldid',
]),
stage_in_argnames=set([
@@ -299,8 +357,26 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
out = jf(
d.qpos.shape[0],
m._impl.block_dim,
m._impl.flex_conaffinity,
m._impl.flex_condim,
m._impl.flex_contype,
m._impl.flex_dim,
m._impl.flex_elem,
m._impl.flex_elemadr,
m._impl.flex_elemdataadr,
m._impl.flex_elemnum,
m._impl.flex_friction,
m._impl.flex_margin,
m._impl.flex_radius,
m._impl.flex_shell,
m._impl.flex_shelldataadr,
m._impl.flex_shellnum,
m._impl.flex_vertadr,
m._impl.flex_vertflexid,
m.geom_aabb,
m.geom_conaffinity,
m.geom_condim,
m.geom_contype,
m.geom_dataid,
m.geom_friction,
m.geom_gap,
@@ -337,6 +413,10 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.mesh_vert,
m.mesh_vertadr,
m.mesh_vertnum,
m._impl.nflex,
m._impl.nflexelem,
m._impl.nflexshelldata,
m._impl.nflexvert,
m.ngeom,
m._impl.nmaxmeshdeg,
m._impl.nmaxpolygon,
@@ -366,12 +446,15 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.opt._impl.sdf_iterations,
d._impl.naccdmax,
d._impl.naconmax,
d._impl.flexvert_xpos,
d.geom_xmat,
d.geom_xpos,
d._impl.nacon,
d._impl.ncollision,
d._impl.contact__dim,
d._impl.contact__dist,
d._impl.contact__efc_address,
d._impl.contact__flex,
d._impl.contact__frame,
d._impl.contact__friction,
d._impl.contact__geom,
@@ -382,6 +465,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
d._impl.contact__solref,
d._impl.contact__solreffriction,
d._impl.contact__type,
d._impl.contact__vert,
d._impl.contact__worldid,
)
d = d.tree_replace({
@@ -389,17 +473,20 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'_impl.ncollision': out[1],
'_impl.contact__dim': out[2],
'_impl.contact__dist': out[3],
'_impl.contact__frame': out[4],
'_impl.contact__friction': out[5],
'_impl.contact__geom': out[6],
'_impl.contact__geomcollisionid': out[7],
'_impl.contact__includemargin': out[8],
'_impl.contact__pos': out[9],
'_impl.contact__solimp': out[10],
'_impl.contact__solref': out[11],
'_impl.contact__solreffriction': out[12],
'_impl.contact__type': out[13],
'_impl.contact__worldid': out[14],
'_impl.contact__efc_address': out[4],
'_impl.contact__flex': out[5],
'_impl.contact__frame': out[6],
'_impl.contact__friction': out[7],
'_impl.contact__geom': out[8],
'_impl.contact__geomcollisionid': out[9],
'_impl.contact__includemargin': out[10],
'_impl.contact__pos': out[11],
'_impl.contact__solimp': out[12],
'_impl.contact__solref': out[13],
'_impl.contact__solreffriction': out[14],
'_impl.contact__type': out[15],
'_impl.contact__vert': out[16],
'_impl.contact__worldid': out[17],
})
return d
mjx/mujoco/mjx/warp/forward.py
+232 -58
View File
@@ -138,6 +138,10 @@ def _forward_shim(
eq_type: wp.array(dtype=int),
eq_wld_adr: wp.array(dtype=int),
flex_bending: wp.array2d(dtype=float),
flex_centered: wp.array(dtype=bool),
flex_conaffinity: wp.array(dtype=int),
flex_condim: wp.array(dtype=int),
flex_contype: wp.array(dtype=int),
flex_damping: wp.array(dtype=float),
flex_dim: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
@@ -146,12 +150,22 @@ def _forward_shim(
flex_edgenum: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elemedge: wp.array(dtype=int),
flex_elemedgeadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_friction: wp.array(dtype=wp.vec3),
flex_margin: wp.array(dtype=float),
flex_radius: wp.array(dtype=float),
flex_shell: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_shellnum: wp.array(dtype=int),
flex_stiffness: wp.array2d(dtype=float),
flex_vert: wp.array(dtype=wp.vec3),
flex_vertadr: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flex_vertflexid: wp.array(dtype=int),
flex_vertnum: wp.array(dtype=int),
flexedge_J_colind: wp.array(dtype=int),
flexedge_J_rowadr: wp.array(dtype=int),
flexedge_J_rownnz: wp.array(dtype=int),
@@ -159,7 +173,9 @@ def _forward_shim(
flexedge_length0: wp.array(dtype=float),
geom_aabb: wp.array3d(dtype=wp.vec3),
geom_bodyid: wp.array(dtype=int),
geom_conaffinity: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_contype: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_fluid: wp.array2d(dtype=float),
geom_friction: wp.array2d(dtype=wp.vec3),
@@ -213,6 +229,7 @@ def _forward_shim(
light_targetbodyid: wp.array(dtype=int),
mapM2M: wp.array(dtype=int),
mat_rgba: wp.array2d(dtype=wp.vec4),
max_ten_J_rownnz: int,
mesh_face: wp.array(dtype=wp.vec3i),
mesh_faceadr: wp.array(dtype=int),
mesh_graph: wp.array(dtype=int),
@@ -235,6 +252,7 @@ def _forward_shim(
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
nC: int,
nJten: int,
na: int,
nacttrnbody: int,
nbody: int,
@@ -244,6 +262,7 @@ def _forward_shim(
nflex: int,
nflexedge: int,
nflexelem: int,
nflexshelldata: int,
nflexvert: int,
ngeom: int,
ngravcomp: int,
@@ -279,7 +298,9 @@ def _forward_shim(
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
qLD_all_updates: wp.array(dtype=wp.vec3i),
qLD_level_offsets: wp.array(dtype=int),
qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...],
qM_fullm_i: wp.array(dtype=int),
qM_fullm_j: wp.array(dtype=int),
@@ -323,6 +344,9 @@ def _forward_shim(
site_type: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int),
taxel_vertadr: wp.array(dtype=int),
ten_J_colind: wp.array(dtype=int),
ten_J_rowadr: wp.array(dtype=int),
ten_J_rownnz: wp.array(dtype=int),
tendon_actfrclimited: wp.array(dtype=bool),
tendon_actfrcrange: wp.array2d(dtype=wp.vec2),
tendon_adr: wp.array(dtype=int),
@@ -382,6 +406,7 @@ def _forward_shim(
naccdmax: int,
naconmax: int,
njmax: int,
njmax_nnz: int,
act: wp.array2d(dtype=float),
act_dot: wp.array2d(dtype=float),
actuator_force: wp.array2d(dtype=float),
@@ -446,7 +471,7 @@ def _forward_shim(
subtree_angmom: wp.array2d(dtype=wp.vec3),
subtree_com: wp.array2d(dtype=wp.vec3),
subtree_linvel: wp.array2d(dtype=wp.vec3),
ten_J: wp.array3d(dtype=float),
ten_J: wp.array2d(dtype=float),
ten_length: wp.array2d(dtype=float),
ten_velocity: wp.array2d(dtype=float),
ten_wrapadr: wp.array2d(dtype=int),
@@ -466,6 +491,7 @@ def _forward_shim(
contact__dim: wp.array(dtype=int),
contact__dist: wp.array(dtype=float),
contact__efc_address: wp.array2d(dtype=int),
contact__flex: wp.array(dtype=wp.vec2i),
contact__frame: wp.array(dtype=wp.mat33),
contact__friction: wp.array(dtype=mjwp_types.vec5),
contact__geom: wp.array(dtype=wp.vec2i),
@@ -476,6 +502,7 @@ def _forward_shim(
contact__solref: wp.array(dtype=wp.vec2),
contact__solreffriction: wp.array(dtype=wp.vec2),
contact__type: wp.array(dtype=int),
contact__vert: wp.array(dtype=wp.vec2i),
contact__worldid: wp.array(dtype=int),
efc__D: wp.array2d(dtype=float),
efc__J: wp.array3d(dtype=float),
@@ -585,6 +612,10 @@ def _forward_shim(
_m.eq_type = eq_type
_m.eq_wld_adr = eq_wld_adr
_m.flex_bending = flex_bending
_m.flex_centered = flex_centered
_m.flex_conaffinity = flex_conaffinity
_m.flex_condim = flex_condim
_m.flex_contype = flex_contype
_m.flex_damping = flex_damping
_m.flex_dim = flex_dim
_m.flex_edge = flex_edge
@@ -593,12 +624,22 @@ def _forward_shim(
_m.flex_edgenum = flex_edgenum
_m.flex_elem = flex_elem
_m.flex_elemadr = flex_elemadr
_m.flex_elemdataadr = flex_elemdataadr
_m.flex_elemedge = flex_elemedge
_m.flex_elemedgeadr = flex_elemedgeadr
_m.flex_elemnum = flex_elemnum
_m.flex_friction = flex_friction
_m.flex_margin = flex_margin
_m.flex_radius = flex_radius
_m.flex_shell = flex_shell
_m.flex_shelldataadr = flex_shelldataadr
_m.flex_shellnum = flex_shellnum
_m.flex_stiffness = flex_stiffness
_m.flex_vert = flex_vert
_m.flex_vertadr = flex_vertadr
_m.flex_vertbodyid = flex_vertbodyid
_m.flex_vertflexid = flex_vertflexid
_m.flex_vertnum = flex_vertnum
_m.flexedge_J_colind = flexedge_J_colind
_m.flexedge_J_rowadr = flexedge_J_rowadr
_m.flexedge_J_rownnz = flexedge_J_rownnz
@@ -606,7 +647,9 @@ def _forward_shim(
_m.flexedge_length0 = flexedge_length0
_m.geom_aabb = geom_aabb
_m.geom_bodyid = geom_bodyid
_m.geom_conaffinity = geom_conaffinity
_m.geom_condim = geom_condim
_m.geom_contype = geom_contype
_m.geom_dataid = geom_dataid
_m.geom_fluid = geom_fluid
_m.geom_friction = geom_friction
@@ -660,6 +703,7 @@ def _forward_shim(
_m.light_targetbodyid = light_targetbodyid
_m.mapM2M = mapM2M
_m.mat_rgba = mat_rgba
_m.max_ten_J_rownnz = max_ten_J_rownnz
_m.mesh_face = mesh_face
_m.mesh_faceadr = mesh_faceadr
_m.mesh_graph = mesh_graph
@@ -682,6 +726,7 @@ def _forward_shim(
_m.mesh_vertadr = mesh_vertadr
_m.mesh_vertnum = mesh_vertnum
_m.nC = nC
_m.nJten = nJten
_m.na = na
_m.nacttrnbody = nacttrnbody
_m.nbody = nbody
@@ -691,6 +736,7 @@ def _forward_shim(
_m.nflex = nflex
_m.nflexedge = nflexedge
_m.nflexelem = nflexelem
_m.nflexshelldata = nflexshelldata
_m.nflexvert = nflexvert
_m.ngeom = ngeom
_m.ngravcomp = ngravcomp
@@ -753,6 +799,8 @@ def _forward_shim(
_m.pair_solreffriction = pair_solreffriction
_m.plugin = plugin
_m.plugin_attr = plugin_attr
_m.qLD_all_updates = qLD_all_updates
_m.qLD_level_offsets = qLD_level_offsets
_m.qLD_updates = qLD_updates
_m.qM_fullm_i = qM_fullm_i
_m.qM_fullm_j = qM_fullm_j
@@ -797,6 +845,9 @@ def _forward_shim(
_m.stat.meaninertia = stat__meaninertia
_m.taxel_sensorid = taxel_sensorid
_m.taxel_vertadr = taxel_vertadr
_m.ten_J_colind = ten_J_colind
_m.ten_J_rowadr = ten_J_rowadr
_m.ten_J_rownnz = ten_J_rownnz
_m.tendon_actfrclimited = tendon_actfrclimited
_m.tendon_actfrcrange = tendon_actfrcrange
_m.tendon_adr = tendon_adr
@@ -842,6 +893,7 @@ def _forward_shim(
_d.contact.dim = contact__dim
_d.contact.dist = contact__dist
_d.contact.efc_address = contact__efc_address
_d.contact.flex = contact__flex
_d.contact.frame = contact__frame
_d.contact.friction = contact__friction
_d.contact.geom = contact__geom
@@ -852,6 +904,7 @@ def _forward_shim(
_d.contact.solref = contact__solref
_d.contact.solreffriction = contact__solreffriction
_d.contact.type = contact__type
_d.contact.vert = contact__vert
_d.contact.worldid = contact__worldid
_d.crb = crb
_d.ctrl = ctrl
@@ -895,6 +948,7 @@ def _forward_shim(
_d.nf = nf
_d.nisland = nisland
_d.njmax = njmax
_d.njmax_nnz = njmax_nnz
_d.nl = nl
_d.qLD = qLD
_d.qLDiagInv = qLDiagInv
@@ -1018,6 +1072,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'contact__dim': d._impl.contact__dim.shape,
'contact__dist': d._impl.contact__dist.shape,
'contact__efc_address': d._impl.contact__efc_address.shape,
'contact__flex': d._impl.contact__flex.shape,
'contact__frame': d._impl.contact__frame.shape,
'contact__friction': d._impl.contact__friction.shape,
'contact__geom': d._impl.contact__geom.shape,
@@ -1028,6 +1083,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'contact__solref': d._impl.contact__solref.shape,
'contact__solreffriction': d._impl.contact__solreffriction.shape,
'contact__type': d._impl.contact__type.shape,
'contact__vert': d._impl.contact__vert.shape,
'contact__worldid': d._impl.contact__worldid.shape,
'efc__D': d._impl.efc__D.shape,
'efc__J': d._impl.efc__J.shape,
@@ -1047,7 +1103,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_forward_shim,
num_outputs=100,
num_outputs=102,
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([
@@ -1125,6 +1181,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'contact__dim',
'contact__dist',
'contact__efc_address',
'contact__flex',
'contact__frame',
'contact__friction',
'contact__geom',
@@ -1135,6 +1192,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'contact__solref',
'contact__solreffriction',
'contact__type',
'contact__vert',
'contact__worldid',
'efc__D',
'efc__J',
@@ -1417,6 +1475,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.eq_type,
m._impl.eq_wld_adr,
m._impl.flex_bending,
m._impl.flex_centered,
m._impl.flex_conaffinity,
m._impl.flex_condim,
m._impl.flex_contype,
m._impl.flex_damping,
m._impl.flex_dim,
m._impl.flex_edge,
@@ -1425,12 +1487,22 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.flex_edgenum,
m._impl.flex_elem,
m._impl.flex_elemadr,
m._impl.flex_elemdataadr,
m._impl.flex_elemedge,
m._impl.flex_elemedgeadr,
m._impl.flex_elemnum,
m._impl.flex_friction,
m._impl.flex_margin,
m._impl.flex_radius,
m._impl.flex_shell,
m._impl.flex_shelldataadr,
m._impl.flex_shellnum,
m._impl.flex_stiffness,
m._impl.flex_vert,
m._impl.flex_vertadr,
m._impl.flex_vertbodyid,
m._impl.flex_vertflexid,
m._impl.flex_vertnum,
m._impl.flexedge_J_colind,
m._impl.flexedge_J_rowadr,
m._impl.flexedge_J_rownnz,
@@ -1438,7 +1510,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.flexedge_length0,
m.geom_aabb,
m.geom_bodyid,
m.geom_conaffinity,
m.geom_condim,
m.geom_contype,
m.geom_dataid,
m.geom_fluid,
m.geom_friction,
@@ -1492,6 +1566,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.light_targetbodyid,
m._impl.mapM2M,
m.mat_rgba,
m._impl.max_ten_J_rownnz,
m.mesh_face,
m.mesh_faceadr,
m.mesh_graph,
@@ -1514,6 +1589,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.mesh_vertadr,
m.mesh_vertnum,
m.nC,
m.nJten,
m.na,
m._impl.nacttrnbody,
m.nbody,
@@ -1523,6 +1599,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.nflex,
m._impl.nflexedge,
m._impl.nflexelem,
m._impl.nflexshelldata,
m._impl.nflexvert,
m.ngeom,
m.ngravcomp,
@@ -1559,6 +1636,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.pair_solreffriction,
m._impl.plugin,
m._impl.plugin_attr,
m._impl.qLD_all_updates,
m._impl.qLD_level_offsets,
m._impl.qLD_updates,
m._impl.qM_fullm_i,
m._impl.qM_fullm_j,
@@ -1602,6 +1681,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.site_type,
m._impl.taxel_sensorid,
m._impl.taxel_vertadr,
m._impl.ten_J_colind,
m._impl.ten_J_rowadr,
m._impl.ten_J_rownnz,
m.tendon_actfrclimited,
m.tendon_actfrcrange,
m.tendon_adr,
@@ -1660,6 +1742,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.naccdmax,
d._impl.naconmax,
d._impl.njmax,
d._impl.njmax_nnz,
d.act,
d.act_dot,
d.actuator_force,
@@ -1744,6 +1827,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.contact__dim,
d._impl.contact__dist,
d._impl.contact__efc_address,
d._impl.contact__flex,
d._impl.contact__frame,
d._impl.contact__friction,
d._impl.contact__geom,
@@ -1754,6 +1838,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.contact__solref,
d._impl.contact__solreffriction,
d._impl.contact__type,
d._impl.contact__vert,
d._impl.contact__worldid,
d._impl.efc__D,
d._impl.efc__J,
@@ -1846,32 +1931,34 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'_impl.contact__dim': out[71],
'_impl.contact__dist': out[72],
'_impl.contact__efc_address': out[73],
'_impl.contact__frame': out[74],
'_impl.contact__friction': out[75],
'_impl.contact__geom': out[76],
'_impl.contact__geomcollisionid': out[77],
'_impl.contact__includemargin': out[78],
'_impl.contact__pos': out[79],
'_impl.contact__solimp': out[80],
'_impl.contact__solref': out[81],
'_impl.contact__solreffriction': out[82],
'_impl.contact__type': out[83],
'_impl.contact__worldid': out[84],
'_impl.efc__D': out[85],
'_impl.efc__J': out[86],
'_impl.efc__J_colind': out[87],
'_impl.efc__J_rowadr': out[88],
'_impl.efc__J_rownnz': out[89],
'_impl.efc__Ma': out[90],
'_impl.efc__aref': out[91],
'_impl.efc__force': out[92],
'_impl.efc__frictionloss': out[93],
'_impl.efc__id': out[94],
'_impl.efc__margin': out[95],
'_impl.efc__pos': out[96],
'_impl.efc__state': out[97],
'_impl.efc__type': out[98],
'_impl.efc__vel': out[99],
'_impl.contact__flex': out[74],
'_impl.contact__frame': out[75],
'_impl.contact__friction': out[76],
'_impl.contact__geom': out[77],
'_impl.contact__geomcollisionid': out[78],
'_impl.contact__includemargin': out[79],
'_impl.contact__pos': out[80],
'_impl.contact__solimp': out[81],
'_impl.contact__solref': out[82],
'_impl.contact__solreffriction': out[83],
'_impl.contact__type': out[84],
'_impl.contact__vert': out[85],
'_impl.contact__worldid': out[86],
'_impl.efc__D': out[87],
'_impl.efc__J': out[88],
'_impl.efc__J_colind': out[89],
'_impl.efc__J_rowadr': out[90],
'_impl.efc__J_rownnz': out[91],
'_impl.efc__Ma': out[92],
'_impl.efc__aref': out[93],
'_impl.efc__force': out[94],
'_impl.efc__frictionloss': out[95],
'_impl.efc__id': out[96],
'_impl.efc__margin': out[97],
'_impl.efc__pos': out[98],
'_impl.efc__state': out[99],
'_impl.efc__type': out[100],
'_impl.efc__vel': out[101],
})
return d
@@ -1980,6 +2067,10 @@ def _step_shim(
eq_type: wp.array(dtype=int),
eq_wld_adr: wp.array(dtype=int),
flex_bending: wp.array2d(dtype=float),
flex_centered: wp.array(dtype=bool),
flex_conaffinity: wp.array(dtype=int),
flex_condim: wp.array(dtype=int),
flex_contype: wp.array(dtype=int),
flex_damping: wp.array(dtype=float),
flex_dim: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
@@ -1988,12 +2079,22 @@ def _step_shim(
flex_edgenum: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemdataadr: wp.array(dtype=int),
flex_elemedge: wp.array(dtype=int),
flex_elemedgeadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_friction: wp.array(dtype=wp.vec3),
flex_margin: wp.array(dtype=float),
flex_radius: wp.array(dtype=float),
flex_shell: wp.array(dtype=int),
flex_shelldataadr: wp.array(dtype=int),
flex_shellnum: wp.array(dtype=int),
flex_stiffness: wp.array2d(dtype=float),
flex_vert: wp.array(dtype=wp.vec3),
flex_vertadr: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flex_vertflexid: wp.array(dtype=int),
flex_vertnum: wp.array(dtype=int),
flexedge_J_colind: wp.array(dtype=int),
flexedge_J_rowadr: wp.array(dtype=int),
flexedge_J_rownnz: wp.array(dtype=int),
@@ -2001,7 +2102,9 @@ def _step_shim(
flexedge_length0: wp.array(dtype=float),
geom_aabb: wp.array3d(dtype=wp.vec3),
geom_bodyid: wp.array(dtype=int),
geom_conaffinity: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_contype: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_fluid: wp.array2d(dtype=float),
geom_friction: wp.array2d(dtype=wp.vec3),
@@ -2055,6 +2158,7 @@ def _step_shim(
light_targetbodyid: wp.array(dtype=int),
mapM2M: wp.array(dtype=int),
mat_rgba: wp.array2d(dtype=wp.vec4),
max_ten_J_rownnz: int,
mesh_face: wp.array(dtype=wp.vec3i),
mesh_faceadr: wp.array(dtype=int),
mesh_graph: wp.array(dtype=int),
@@ -2077,6 +2181,7 @@ def _step_shim(
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
nC: int,
nJten: int,
nM: int,
na: int,
nacttrnbody: int,
@@ -2087,6 +2192,7 @@ def _step_shim(
nflex: int,
nflexedge: int,
nflexelem: int,
nflexshelldata: int,
nflexvert: int,
ngeom: int,
ngravcomp: int,
@@ -2122,7 +2228,9 @@ def _step_shim(
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
qLD_all_updates: wp.array(dtype=wp.vec3i),
qLD_level_offsets: wp.array(dtype=int),
qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...],
qM_fullm_i: wp.array(dtype=int),
qM_fullm_j: wp.array(dtype=int),
@@ -2166,6 +2274,9 @@ def _step_shim(
site_type: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int),
taxel_vertadr: wp.array(dtype=int),
ten_J_colind: wp.array(dtype=int),
ten_J_rowadr: wp.array(dtype=int),
ten_J_rownnz: wp.array(dtype=int),
tendon_actfrclimited: wp.array(dtype=bool),
tendon_actfrcrange: wp.array2d(dtype=wp.vec2),
tendon_adr: wp.array(dtype=int),
@@ -2226,6 +2337,7 @@ def _step_shim(
naccdmax: int,
naconmax: int,
njmax: int,
njmax_nnz: int,
act: wp.array2d(dtype=float),
act_dot: wp.array2d(dtype=float),
actuator_force: wp.array2d(dtype=float),
@@ -2290,7 +2402,7 @@ def _step_shim(
subtree_angmom: wp.array2d(dtype=wp.vec3),
subtree_com: wp.array2d(dtype=wp.vec3),
subtree_linvel: wp.array2d(dtype=wp.vec3),
ten_J: wp.array3d(dtype=float),
ten_J: wp.array2d(dtype=float),
ten_length: wp.array2d(dtype=float),
ten_velocity: wp.array2d(dtype=float),
ten_wrapadr: wp.array2d(dtype=int),
@@ -2310,6 +2422,7 @@ def _step_shim(
contact__dim: wp.array(dtype=int),
contact__dist: wp.array(dtype=float),
contact__efc_address: wp.array2d(dtype=int),
contact__flex: wp.array(dtype=wp.vec2i),
contact__frame: wp.array(dtype=wp.mat33),
contact__friction: wp.array(dtype=mjwp_types.vec5),
contact__geom: wp.array(dtype=wp.vec2i),
@@ -2320,6 +2433,7 @@ def _step_shim(
contact__solref: wp.array(dtype=wp.vec2),
contact__solreffriction: wp.array(dtype=wp.vec2),
contact__type: wp.array(dtype=int),
contact__vert: wp.array(dtype=wp.vec2i),
contact__worldid: wp.array(dtype=int),
efc__D: wp.array2d(dtype=float),
efc__J: wp.array3d(dtype=float),
@@ -2429,6 +2543,10 @@ def _step_shim(
_m.eq_type = eq_type
_m.eq_wld_adr = eq_wld_adr
_m.flex_bending = flex_bending
_m.flex_centered = flex_centered
_m.flex_conaffinity = flex_conaffinity
_m.flex_condim = flex_condim
_m.flex_contype = flex_contype
_m.flex_damping = flex_damping
_m.flex_dim = flex_dim
_m.flex_edge = flex_edge
@@ -2437,12 +2555,22 @@ def _step_shim(
_m.flex_edgenum = flex_edgenum
_m.flex_elem = flex_elem
_m.flex_elemadr = flex_elemadr
_m.flex_elemdataadr = flex_elemdataadr
_m.flex_elemedge = flex_elemedge
_m.flex_elemedgeadr = flex_elemedgeadr
_m.flex_elemnum = flex_elemnum
_m.flex_friction = flex_friction
_m.flex_margin = flex_margin
_m.flex_radius = flex_radius
_m.flex_shell = flex_shell
_m.flex_shelldataadr = flex_shelldataadr
_m.flex_shellnum = flex_shellnum
_m.flex_stiffness = flex_stiffness
_m.flex_vert = flex_vert
_m.flex_vertadr = flex_vertadr
_m.flex_vertbodyid = flex_vertbodyid
_m.flex_vertflexid = flex_vertflexid
_m.flex_vertnum = flex_vertnum
_m.flexedge_J_colind = flexedge_J_colind
_m.flexedge_J_rowadr = flexedge_J_rowadr
_m.flexedge_J_rownnz = flexedge_J_rownnz
@@ -2450,7 +2578,9 @@ def _step_shim(
_m.flexedge_length0 = flexedge_length0
_m.geom_aabb = geom_aabb
_m.geom_bodyid = geom_bodyid
_m.geom_conaffinity = geom_conaffinity
_m.geom_condim = geom_condim
_m.geom_contype = geom_contype
_m.geom_dataid = geom_dataid
_m.geom_fluid = geom_fluid
_m.geom_friction = geom_friction
@@ -2504,6 +2634,7 @@ def _step_shim(
_m.light_targetbodyid = light_targetbodyid
_m.mapM2M = mapM2M
_m.mat_rgba = mat_rgba
_m.max_ten_J_rownnz = max_ten_J_rownnz
_m.mesh_face = mesh_face
_m.mesh_faceadr = mesh_faceadr
_m.mesh_graph = mesh_graph
@@ -2526,6 +2657,7 @@ def _step_shim(
_m.mesh_vertadr = mesh_vertadr
_m.mesh_vertnum = mesh_vertnum
_m.nC = nC
_m.nJten = nJten
_m.nM = nM
_m.na = na
_m.nacttrnbody = nacttrnbody
@@ -2536,6 +2668,7 @@ def _step_shim(
_m.nflex = nflex
_m.nflexedge = nflexedge
_m.nflexelem = nflexelem
_m.nflexshelldata = nflexshelldata
_m.nflexvert = nflexvert
_m.ngeom = ngeom
_m.ngravcomp = ngravcomp
@@ -2599,6 +2732,8 @@ def _step_shim(
_m.pair_solreffriction = pair_solreffriction
_m.plugin = plugin
_m.plugin_attr = plugin_attr
_m.qLD_all_updates = qLD_all_updates
_m.qLD_level_offsets = qLD_level_offsets
_m.qLD_updates = qLD_updates
_m.qM_fullm_i = qM_fullm_i
_m.qM_fullm_j = qM_fullm_j
@@ -2643,6 +2778,9 @@ def _step_shim(
_m.stat.meaninertia = stat__meaninertia
_m.taxel_sensorid = taxel_sensorid
_m.taxel_vertadr = taxel_vertadr
_m.ten_J_colind = ten_J_colind
_m.ten_J_rowadr = ten_J_rowadr
_m.ten_J_rownnz = ten_J_rownnz
_m.tendon_actfrclimited = tendon_actfrclimited
_m.tendon_actfrcrange = tendon_actfrcrange
_m.tendon_adr = tendon_adr
@@ -2688,6 +2826,7 @@ def _step_shim(
_d.contact.dim = contact__dim
_d.contact.dist = contact__dist
_d.contact.efc_address = contact__efc_address
_d.contact.flex = contact__flex
_d.contact.frame = contact__frame
_d.contact.friction = contact__friction
_d.contact.geom = contact__geom
@@ -2698,6 +2837,7 @@ def _step_shim(
_d.contact.solref = contact__solref
_d.contact.solreffriction = contact__solreffriction
_d.contact.type = contact__type
_d.contact.vert = contact__vert
_d.contact.worldid = contact__worldid
_d.crb = crb
_d.ctrl = ctrl
@@ -2741,6 +2881,7 @@ def _step_shim(
_d.nf = nf
_d.nisland = nisland
_d.njmax = njmax
_d.njmax_nnz = njmax_nnz
_d.nl = nl
_d.qLD = qLD
_d.qLDiagInv = qLDiagInv
@@ -2868,6 +3009,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'contact__dim': d._impl.contact__dim.shape,
'contact__dist': d._impl.contact__dist.shape,
'contact__efc_address': d._impl.contact__efc_address.shape,
'contact__flex': d._impl.contact__flex.shape,
'contact__frame': d._impl.contact__frame.shape,
'contact__friction': d._impl.contact__friction.shape,
'contact__geom': d._impl.contact__geom.shape,
@@ -2878,6 +3020,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'contact__solref': d._impl.contact__solref.shape,
'contact__solreffriction': d._impl.contact__solreffriction.shape,
'contact__type': d._impl.contact__type.shape,
'contact__vert': d._impl.contact__vert.shape,
'contact__worldid': d._impl.contact__worldid.shape,
'efc__D': d._impl.efc__D.shape,
'efc__J': d._impl.efc__J.shape,
@@ -2897,7 +3040,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_step_shim,
num_outputs=104,
num_outputs=106,
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([
@@ -2979,6 +3122,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'contact__dim',
'contact__dist',
'contact__efc_address',
'contact__flex',
'contact__frame',
'contact__friction',
'contact__geom',
@@ -2989,6 +3133,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'contact__solref',
'contact__solreffriction',
'contact__type',
'contact__vert',
'contact__worldid',
'efc__D',
'efc__J',
@@ -3275,6 +3420,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.eq_type,
m._impl.eq_wld_adr,
m._impl.flex_bending,
m._impl.flex_centered,
m._impl.flex_conaffinity,
m._impl.flex_condim,
m._impl.flex_contype,
m._impl.flex_damping,
m._impl.flex_dim,
m._impl.flex_edge,
@@ -3283,12 +3432,22 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.flex_edgenum,
m._impl.flex_elem,
m._impl.flex_elemadr,
m._impl.flex_elemdataadr,
m._impl.flex_elemedge,
m._impl.flex_elemedgeadr,
m._impl.flex_elemnum,
m._impl.flex_friction,
m._impl.flex_margin,
m._impl.flex_radius,
m._impl.flex_shell,
m._impl.flex_shelldataadr,
m._impl.flex_shellnum,
m._impl.flex_stiffness,
m._impl.flex_vert,
m._impl.flex_vertadr,
m._impl.flex_vertbodyid,
m._impl.flex_vertflexid,
m._impl.flex_vertnum,
m._impl.flexedge_J_colind,
m._impl.flexedge_J_rowadr,
m._impl.flexedge_J_rownnz,
@@ -3296,7 +3455,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.flexedge_length0,
m.geom_aabb,
m.geom_bodyid,
m.geom_conaffinity,
m.geom_condim,
m.geom_contype,
m.geom_dataid,
m.geom_fluid,
m.geom_friction,
@@ -3350,6 +3511,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.light_targetbodyid,
m._impl.mapM2M,
m.mat_rgba,
m._impl.max_ten_J_rownnz,
m.mesh_face,
m.mesh_faceadr,
m.mesh_graph,
@@ -3372,6 +3534,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.mesh_vertadr,
m.mesh_vertnum,
m.nC,
m.nJten,
m.nM,
m.na,
m._impl.nacttrnbody,
@@ -3382,6 +3545,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.nflex,
m._impl.nflexedge,
m._impl.nflexelem,
m._impl.nflexshelldata,
m._impl.nflexvert,
m.ngeom,
m.ngravcomp,
@@ -3418,6 +3582,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.pair_solreffriction,
m._impl.plugin,
m._impl.plugin_attr,
m._impl.qLD_all_updates,
m._impl.qLD_level_offsets,
m._impl.qLD_updates,
m._impl.qM_fullm_i,
m._impl.qM_fullm_j,
@@ -3461,6 +3627,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.site_type,
m._impl.taxel_sensorid,
m._impl.taxel_vertadr,
m._impl.ten_J_colind,
m._impl.ten_J_rowadr,
m._impl.ten_J_rownnz,
m.tendon_actfrclimited,
m.tendon_actfrcrange,
m.tendon_adr,
@@ -3520,6 +3689,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.naccdmax,
d._impl.naconmax,
d._impl.njmax,
d._impl.njmax_nnz,
d.act,
d.act_dot,
d.actuator_force,
@@ -3604,6 +3774,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.contact__dim,
d._impl.contact__dist,
d._impl.contact__efc_address,
d._impl.contact__flex,
d._impl.contact__frame,
d._impl.contact__friction,
d._impl.contact__geom,
@@ -3614,6 +3785,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.contact__solref,
d._impl.contact__solreffriction,
d._impl.contact__type,
d._impl.contact__vert,
d._impl.contact__worldid,
d._impl.efc__D,
d._impl.efc__J,
@@ -3710,32 +3882,34 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'_impl.contact__dim': out[75],
'_impl.contact__dist': out[76],
'_impl.contact__efc_address': out[77],
'_impl.contact__frame': out[78],
'_impl.contact__friction': out[79],
'_impl.contact__geom': out[80],
'_impl.contact__geomcollisionid': out[81],
'_impl.contact__includemargin': out[82],
'_impl.contact__pos': out[83],
'_impl.contact__solimp': out[84],
'_impl.contact__solref': out[85],
'_impl.contact__solreffriction': out[86],
'_impl.contact__type': out[87],
'_impl.contact__worldid': out[88],
'_impl.efc__D': out[89],
'_impl.efc__J': out[90],
'_impl.efc__J_colind': out[91],
'_impl.efc__J_rowadr': out[92],
'_impl.efc__J_rownnz': out[93],
'_impl.efc__Ma': out[94],
'_impl.efc__aref': out[95],
'_impl.efc__force': out[96],
'_impl.efc__frictionloss': out[97],
'_impl.efc__id': out[98],
'_impl.efc__margin': out[99],
'_impl.efc__pos': out[100],
'_impl.efc__state': out[101],
'_impl.efc__type': out[102],
'_impl.efc__vel': out[103],
'_impl.contact__flex': out[78],
'_impl.contact__frame': out[79],
'_impl.contact__friction': out[80],
'_impl.contact__geom': out[81],
'_impl.contact__geomcollisionid': out[82],
'_impl.contact__includemargin': out[83],
'_impl.contact__pos': out[84],
'_impl.contact__solimp': out[85],
'_impl.contact__solref': out[86],
'_impl.contact__solreffriction': out[87],
'_impl.contact__type': out[88],
'_impl.contact__vert': out[89],
'_impl.contact__worldid': out[90],
'_impl.efc__D': out[91],
'_impl.efc__J': out[92],
'_impl.efc__J_colind': out[93],
'_impl.efc__J_rowadr': out[94],
'_impl.efc__J_rownnz': out[95],
'_impl.efc__Ma': out[96],
'_impl.efc__aref': out[97],
'_impl.efc__force': out[98],
'_impl.efc__frictionloss': out[99],
'_impl.efc__id': out[100],
'_impl.efc__margin': out[101],
'_impl.efc__pos': out[102],
'_impl.efc__state': out[103],
'_impl.efc__type': out[104],
'_impl.efc__vel': out[105],
})
return d
mjx/mujoco/mjx/warp/forward_test.py
+10 -1
View File
@@ -157,7 +157,16 @@ class ForwardTest(parameterized.TestCase):
m.ten_J_rowadr,
m.ten_J_colind,
)
tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J')
# convert sparse warp ten_J to dense representation
warp_ten_J = np.zeros((m.ntendon, m.nv))
mujoco.mju_sparse2dense(
warp_ten_J,
np.asarray(dx._impl.ten_J),
mx._impl.ten_J_rownnz,
mx._impl.ten_J_rowadr,
mx._impl.ten_J_colind,
)
tu.assert_eq(warp_ten_J, ten_J, 'ten_J')
tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')
tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum')
tu.assert_attr_eq(dx._impl, d, 'wrap_xpos')
mjx/mujoco/mjx/warp/render.py
+13 -3
View File
@@ -48,6 +48,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _render_shim(
# Model
@@ -56,6 +57,9 @@ def _render_shim(
cam_intrinsic: wp.array2d(dtype=wp.vec4),
cam_projection: wp.array(dtype=int),
cam_sensorsize: wp.array(dtype=wp.vec2),
flex_edge: wp.array(dtype=wp.vec2i),
flex_radius: wp.array(dtype=float),
flex_vertadr: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_matid: wp.array2d(dtype=int),
geom_rgba: wp.array2d(dtype=wp.vec4),
@@ -68,11 +72,11 @@ def _render_shim(
mat_texid: wp.array3d(dtype=int),
mat_texrepeat: wp.array2d(dtype=wp.vec2),
mesh_faceadr: wp.array(dtype=int),
nflex: int,
nlight: int,
# Data
cam_xmat: wp.array2d(dtype=wp.mat33),
cam_xpos: wp.array2d(dtype=wp.vec3),
flexvert_xpos: wp.array2d(dtype=wp.vec3),
geom_xmat: wp.array2d(dtype=wp.mat33),
geom_xpos: wp.array2d(dtype=wp.vec3),
light_xdir: wp.array2d(dtype=wp.vec3),
@@ -91,6 +95,9 @@ def _render_shim(
_m.cam_intrinsic = cam_intrinsic
_m.cam_projection = cam_projection
_m.cam_sensorsize = cam_sensorsize
_m.flex_edge = flex_edge
_m.flex_radius = flex_radius
_m.flex_vertadr = flex_vertadr
_m.geom_dataid = geom_dataid
_m.geom_matid = geom_matid
_m.geom_rgba = geom_rgba
@@ -103,10 +110,10 @@ def _render_shim(
_m.mat_texid = mat_texid
_m.mat_texrepeat = mat_texrepeat
_m.mesh_faceadr = mesh_faceadr
_m.nflex = nflex
_m.nlight = nlight
_d.cam_xmat = cam_xmat
_d.cam_xpos = cam_xpos
_d.flexvert_xpos = flexvert_xpos
_d.geom_xmat = geom_xmat
_d.geom_xpos = geom_xpos
_d.light_xdir = light_xdir
@@ -155,6 +162,9 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
m.cam_intrinsic,
m._impl.cam_projection,
m.cam_sensorsize,
m._impl.flex_edge,
m._impl.flex_radius,
m._impl.flex_vertadr,
m.geom_dataid,
m.geom_matid,
m.geom_rgba,
@@ -167,10 +177,10 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
m.mat_texid,
m._impl.mat_texrepeat,
m.mesh_faceadr,
m._impl.nflex,
m.nlight,
d.cam_xmat,
d.cam_xpos,
d._impl.flexvert_xpos,
d.geom_xmat,
d.geom_xpos,
d._impl.light_xdir,
mjx/mujoco/mjx/warp/smooth.py
+16 -7
View File
@@ -297,17 +297,20 @@ def kinematics_vmap(
def _tendon_shim(
# Model
nworld: int,
body_dofadr: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int),
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
dof_bodyid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
jnt_dofadr: wp.array(dtype=int),
jnt_qposadr: wp.array(dtype=int),
ntendon: int,
nv: int,
nwrap: int,
site_bodyid: wp.array(dtype=int),
ten_J_colind: wp.array(dtype=int),
ten_J_rowadr: wp.array(dtype=int),
ten_J_rownnz: wp.array(dtype=int),
tendon_adr: wp.array(dtype=int),
tendon_geom_adr: wp.array(dtype=int),
tendon_jnt_adr: wp.array(dtype=int),
@@ -327,7 +330,7 @@ def _tendon_shim(
qpos: wp.array2d(dtype=float),
site_xpos: wp.array2d(dtype=wp.vec3),
subtree_com: wp.array2d(dtype=wp.vec3),
ten_J: wp.array3d(dtype=float),
ten_J: wp.array2d(dtype=float),
ten_length: wp.array2d(dtype=float),
ten_wrapadr: wp.array2d(dtype=int),
ten_wrapnum: wp.array2d(dtype=int),
@@ -339,17 +342,20 @@ def _tendon_shim(
_m.callback = _cb
_d.efc = _e
_d.contact = _c
_m.body_dofadr = body_dofadr
_m.body_dofnum = body_dofnum
_m.body_parentid = body_parentid
_m.body_rootid = body_rootid
_m.dof_bodyid = dof_bodyid
_m.geom_bodyid = geom_bodyid
_m.geom_size = geom_size
_m.jnt_dofadr = jnt_dofadr
_m.jnt_qposadr = jnt_qposadr
_m.ntendon = ntendon
_m.nv = nv
_m.nwrap = nwrap
_m.site_bodyid = site_bodyid
_m.ten_J_colind = ten_J_colind
_m.ten_J_rowadr = ten_J_rowadr
_m.ten_J_rownnz = ten_J_rownnz
_m.tendon_adr = tendon_adr
_m.tendon_geom_adr = tendon_geom_adr
_m.tendon_jnt_adr = tendon_jnt_adr
@@ -416,17 +422,20 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
)
out = jf(
d.qpos.shape[0],
m.body_dofadr,
m.body_dofnum,
m.body_parentid,
m.body_rootid,
m.dof_bodyid,
m.geom_bodyid,
m.geom_size,
m.jnt_dofadr,
m.jnt_qposadr,
m.ntendon,
m.nv,
m.nwrap,
m.site_bodyid,
m._impl.ten_J_colind,
m._impl.ten_J_rowadr,
m._impl.ten_J_rownnz,
m.tendon_adr,
m._impl.tendon_geom_adr,
m._impl.tendon_jnt_adr,
mjx/mujoco/mjx/warp/types.py
+78 -4
View File
@@ -83,7 +83,7 @@ class BlockDim:
qderiv_actuator_dense: int
ray: int
segmented_sort: int
tendon_velocity: int
solve_LD_sparse_fused: int
update_gradient_JTDAJ_dense: int
update_gradient_JTDAJ_sparse: int
update_gradient_cholesky: int
@@ -141,6 +141,10 @@ class ModelWarp(PyTreeNode):
eq_ten_adr: np.ndarray
eq_wld_adr: np.ndarray
flex_bending: np.ndarray
flex_centered: np.ndarray
flex_conaffinity: np.ndarray
flex_condim: np.ndarray
flex_contype: np.ndarray
flex_damping: np.ndarray
flex_dim: np.ndarray
flex_edge: np.ndarray
@@ -149,12 +153,21 @@ class ModelWarp(PyTreeNode):
flex_edgenum: np.ndarray
flex_elem: np.ndarray
flex_elemadr: np.ndarray
flex_elemdataadr: np.ndarray
flex_elemedge: np.ndarray
flex_elemedgeadr: np.ndarray
flex_elemnum: np.ndarray
flex_friction: np.ndarray
flex_margin: np.ndarray
flex_radius: np.ndarray
flex_shell: np.ndarray
flex_shelldataadr: np.ndarray
flex_shellnum: np.ndarray
flex_stiffness: np.ndarray
flex_vert: np.ndarray
flex_vertadr: np.ndarray
flex_vertbodyid: np.ndarray
flex_vertflexid: np.ndarray
flex_vertnum: np.ndarray
flexedge_J_colind: np.ndarray
flexedge_J_rowadr: np.ndarray
@@ -173,6 +186,7 @@ class ModelWarp(PyTreeNode):
light_targetbodyid: np.ndarray
mapM2M: np.ndarray
mat_texrepeat: jax.Array
max_ten_J_rownnz: int
mesh_polyadr: np.ndarray
mesh_polymap: np.ndarray
mesh_polymapadr: np.ndarray
@@ -191,6 +205,7 @@ class ModelWarp(PyTreeNode):
nflexelem: int
nflexelemdata: int
nflexelemedge: int
nflexshelldata: int
nflexvert: int
nmaxcondim: int
nmaxmeshdeg: int
@@ -213,6 +228,8 @@ class ModelWarp(PyTreeNode):
oct_coeff: np.ndarray
plugin: np.ndarray
plugin_attr: np.ndarray
qLD_all_updates: np.ndarray
qLD_level_offsets: np.ndarray
qLD_updates: Tuple[np.ndarray, ...]
qM_fullm_i: np.ndarray
qM_fullm_j: np.ndarray
@@ -240,6 +257,9 @@ class ModelWarp(PyTreeNode):
sensor_vel_adr: np.ndarray
taxel_sensorid: np.ndarray
taxel_vertadr: np.ndarray
ten_J_colind: np.ndarray
ten_J_rowadr: np.ndarray
ten_J_rownnz: np.ndarray
ten_wrapadr_site: np.ndarray
ten_wrapnum_site: np.ndarray
tendon_geom_adr: np.ndarray
@@ -266,6 +286,7 @@ class DataWarp(PyTreeNode):
contact__dim: jax.Array
contact__dist: jax.Array
contact__efc_address: jax.Array
contact__flex: jax.Array
contact__frame: jax.Array
contact__friction: jax.Array
contact__geom: jax.Array
@@ -276,6 +297,7 @@ class DataWarp(PyTreeNode):
contact__solref: jax.Array
contact__solreffriction: jax.Array
contact__type: jax.Array
contact__vert: jax.Array
contact__worldid: jax.Array
crb: jax.Array
efc__D: jax.Array
@@ -312,6 +334,7 @@ class DataWarp(PyTreeNode):
nf: jax.Array
nisland: jax.Array
njmax: int
njmax_nnz: int
njmax_pad: int
nl: jax.Array
nworld: int
@@ -335,6 +358,7 @@ DATA_NON_VMAP = {
'contact__dim',
'contact__dist',
'contact__efc_address',
'contact__flex',
'contact__frame',
'contact__friction',
'contact__geom',
@@ -345,12 +369,14 @@ DATA_NON_VMAP = {
'contact__solref',
'contact__solreffriction',
'contact__type',
'contact__vert',
'contact__worldid',
'naccdmax',
'nacon',
'naconmax',
'ncollision',
'njmax',
'njmax_nnz',
'njmax_pad',
'nworld',
}
@@ -398,6 +424,7 @@ _NDIM = {
'contact__dim': 1,
'contact__dist': 1,
'contact__efc_address': 2,
'contact__flex': 2,
'contact__frame': 3,
'contact__friction': 2,
'contact__geom': 2,
@@ -408,6 +435,7 @@ _NDIM = {
'contact__solref': 2,
'contact__solreffriction': 2,
'contact__type': 1,
'contact__vert': 2,
'contact__worldid': 1,
'crb': 3,
'ctrl': 2,
@@ -451,6 +479,7 @@ _NDIM = {
'nf': 1,
'nisland': 1,
'njmax': 0,
'njmax_nnz': 0,
'njmax_pad': 0,
'nl': 1,
'nworld': 0,
@@ -480,7 +509,7 @@ _NDIM = {
'subtree_angmom': 3,
'subtree_com': 3,
'subtree_linvel': 3,
'ten_J': 3,
'ten_J': 2,
'ten_length': 2,
'ten_velocity': 2,
'ten_wrapadr': 2,
@@ -535,7 +564,7 @@ _NDIM = {
'block_dim__qderiv_actuator_dense': 0,
'block_dim__ray': 0,
'block_dim__segmented_sort': 0,
'block_dim__tendon_velocity': 0,
'block_dim__solve_LD_sparse_fused': 0,
'block_dim__update_gradient_JTDAJ_dense': 0,
'block_dim__update_gradient_JTDAJ_sparse': 0,
'block_dim__update_gradient_cholesky': 0,
@@ -608,6 +637,10 @@ _NDIM = {
'eq_wld_adr': 1,
'exclude_signature': 1,
'flex_bending': 2,
'flex_centered': 1,
'flex_conaffinity': 1,
'flex_condim': 1,
'flex_contype': 1,
'flex_damping': 1,
'flex_dim': 1,
'flex_edge': 2,
@@ -616,12 +649,21 @@ _NDIM = {
'flex_edgenum': 1,
'flex_elem': 1,
'flex_elemadr': 1,
'flex_elemdataadr': 1,
'flex_elemedge': 1,
'flex_elemedgeadr': 1,
'flex_elemnum': 1,
'flex_friction': 2,
'flex_margin': 1,
'flex_radius': 1,
'flex_shell': 1,
'flex_shelldataadr': 1,
'flex_shellnum': 1,
'flex_stiffness': 2,
'flex_vert': 2,
'flex_vertadr': 1,
'flex_vertbodyid': 1,
'flex_vertflexid': 1,
'flex_vertnum': 1,
'flexedge_J_colind': 1,
'flexedge_J_rowadr': 1,
@@ -692,6 +734,7 @@ _NDIM = {
'mat_rgba': 3,
'mat_texid': 3,
'mat_texrepeat': 3,
'max_ten_J_rownnz': 0,
'mesh_face': 2,
'mesh_faceadr': 1,
'mesh_graph': 1,
@@ -717,6 +760,7 @@ _NDIM = {
'nC': 0,
'nJfe': 0,
'nJmom': 0,
'nJten': 0,
'nM': 0,
'na': 0,
'nacttrnbody': 0,
@@ -730,6 +774,7 @@ _NDIM = {
'nflexelem': 0,
'nflexelemdata': 0,
'nflexelemedge': 0,
'nflexshelldata': 0,
'nflexvert': 0,
'ngeom': 0,
'ngravcomp': 0,
@@ -811,6 +856,8 @@ _NDIM = {
'pair_solreffriction': 3,
'plugin': 1,
'plugin_attr': 2,
'qLD_all_updates': 2,
'qLD_level_offsets': 1,
'qLD_updates': -1,
'qM_fullm_i': 1,
'qM_fullm_j': 1,
@@ -856,6 +903,9 @@ _NDIM = {
'stat__meaninertia': 1,
'taxel_sensorid': 1,
'taxel_vertadr': 1,
'ten_J_colind': 1,
'ten_J_rowadr': 1,
'ten_J_rownnz': 1,
'ten_wrapadr_site': 1,
'ten_wrapnum_site': 1,
'tendon_actfrclimited': 1,
@@ -940,6 +990,7 @@ _BATCH_DIM = {
'contact__dim': False,
'contact__dist': False,
'contact__efc_address': False,
'contact__flex': False,
'contact__frame': False,
'contact__friction': False,
'contact__geom': False,
@@ -950,6 +1001,7 @@ _BATCH_DIM = {
'contact__solref': False,
'contact__solreffriction': False,
'contact__type': False,
'contact__vert': False,
'contact__worldid': False,
'crb': True,
'ctrl': True,
@@ -993,6 +1045,7 @@ _BATCH_DIM = {
'nf': True,
'nisland': True,
'njmax': False,
'njmax_nnz': False,
'njmax_pad': False,
'nl': True,
'nworld': False,
@@ -1077,7 +1130,7 @@ _BATCH_DIM = {
'block_dim__qderiv_actuator_dense': False,
'block_dim__ray': False,
'block_dim__segmented_sort': False,
'block_dim__tendon_velocity': False,
'block_dim__solve_LD_sparse_fused': False,
'block_dim__update_gradient_JTDAJ_dense': False,
'block_dim__update_gradient_JTDAJ_sparse': False,
'block_dim__update_gradient_cholesky': False,
@@ -1150,6 +1203,10 @@ _BATCH_DIM = {
'eq_wld_adr': False,
'exclude_signature': False,
'flex_bending': False,
'flex_centered': False,
'flex_conaffinity': False,
'flex_condim': False,
'flex_contype': False,
'flex_damping': False,
'flex_dim': False,
'flex_edge': False,
@@ -1158,12 +1215,21 @@ _BATCH_DIM = {
'flex_edgenum': False,
'flex_elem': False,
'flex_elemadr': False,
'flex_elemdataadr': False,
'flex_elemedge': False,
'flex_elemedgeadr': False,
'flex_elemnum': False,
'flex_friction': False,
'flex_margin': False,
'flex_radius': False,
'flex_shell': False,
'flex_shelldataadr': False,
'flex_shellnum': False,
'flex_stiffness': False,
'flex_vert': False,
'flex_vertadr': False,
'flex_vertbodyid': False,
'flex_vertflexid': False,
'flex_vertnum': False,
'flexedge_J_colind': False,
'flexedge_J_rowadr': False,
@@ -1234,6 +1300,7 @@ _BATCH_DIM = {
'mat_rgba': True,
'mat_texid': True,
'mat_texrepeat': True,
'max_ten_J_rownnz': False,
'mesh_face': False,
'mesh_faceadr': False,
'mesh_graph': False,
@@ -1259,6 +1326,7 @@ _BATCH_DIM = {
'nC': False,
'nJfe': False,
'nJmom': False,
'nJten': False,
'nM': False,
'na': False,
'nacttrnbody': False,
@@ -1272,6 +1340,7 @@ _BATCH_DIM = {
'nflexelem': False,
'nflexelemdata': False,
'nflexelemedge': False,
'nflexshelldata': False,
'nflexvert': False,
'ngeom': False,
'ngravcomp': False,
@@ -1353,6 +1422,8 @@ _BATCH_DIM = {
'pair_solreffriction': True,
'plugin': False,
'plugin_attr': False,
'qLD_all_updates': False,
'qLD_level_offsets': False,
'qLD_updates': False,
'qM_fullm_i': False,
'qM_fullm_j': False,
@@ -1398,6 +1469,9 @@ _BATCH_DIM = {
'stat__meaninertia': True,
'taxel_sensorid': False,
'taxel_vertadr': False,
'ten_J_colind': False,
'ten_J_rowadr': False,
'ten_J_rownnz': False,
'ten_wrapadr_site': False,
'ten_wrapnum_site': False,
'tendon_actfrclimited': False,