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: 899731493
Change-Id: Ib563023047e85db1aff56e78724dc28d769ef91e
This commit is contained in:
Taylor Howell
2026-04-14 12:55:06 -07:00
committed by Copybara-Service
parent 5724158593
commit 8bf11166b1
34 changed files with 4556 additions and 4468 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+3
View File
@@ -456,6 +456,9 @@ def _put_model_warp(
if not hasattr(mw, k) or k in ('stat', 'opt'): if not hasattr(mw, k) or k in ('stat', 'opt'):
continue continue
field = _wp_to_np_type(getattr(mw, k), k) field = _wp_to_np_type(getattr(mw, k), k)
if k == 'geom_dataid' and field.ndim > 1:
# Batched geom_dataid is not supported in MJX.
field = field[0]
fields[k] = field fields[k] = field
impl_fields = {} impl_fields = {}
mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
Vendored
+6 -6
View File
@@ -41,18 +41,18 @@ def _sum(stack1, stack2):
@wp.kernel @wp.kernel
def ctrl_noise( def ctrl_noise(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
actuator_ctrllimited: wp.array(dtype=bool), actuator_ctrllimited: wp.array[bool],
actuator_ctrlrange: wp.array2d(dtype=wp.vec2), actuator_ctrlrange: wp.array2d[wp.vec2],
# Data in: # Data in:
ctrl_in: wp.array2d(dtype=float), ctrl_in: wp.array2d[float],
# In: # In:
ctrl_center: wp.array1d(dtype=float), ctrl_center: wp.array[float],
step: int, step: int,
ctrlnoisestd: float, ctrlnoisestd: float,
ctrlnoiserate: float, ctrlnoiserate: float,
# Data out: # Data out:
ctrl_out: wp.array2d(dtype=float), ctrl_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
Vendored
+5 -5
View File
@@ -23,10 +23,10 @@ def create_blocked_cholesky_func(block_size: int):
@wp.func @wp.func
def blocked_cholesky_func( def blocked_cholesky_func(
# In: # In:
A: wp.array(dtype=float, ndim=2), A: wp.array2d[float],
matrix_size: int, matrix_size: int,
# Out: # Out:
L: wp.array(dtype=float, ndim=2), L: wp.array2d[float],
): ):
"""Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks. """Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
@@ -68,11 +68,11 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
@wp.func @wp.func
def blocked_cholesky_solve_func( def blocked_cholesky_solve_func(
# In: # In:
L: wp.array(dtype=float, ndim=2), L: wp.array2d[float],
b: wp.array(dtype=float, ndim=2), b: wp.array2d[float],
matrix_size: int, matrix_size: int,
# Out: # Out:
x: wp.array(dtype=float, ndim=2), x: wp.array2d[float],
): ):
"""Block Cholesky factorization and solve. """Block Cholesky factorization and solve.
mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
Vendored
+76 -66
View File
@@ -173,21 +173,21 @@ def _compute_cylinder_bounds(
@wp.kernel @wp.kernel
def _compute_bvh_bounds( def _compute_bvh_bounds(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
bvh_ngeom: int, bvh_ngeom: int,
enabled_geom_ids: wp.array(dtype=int), enabled_geom_ids: wp.array[int],
mesh_bounds_size: wp.array(dtype=wp.vec3), mesh_bounds_size: wp.array[wp.vec3],
hfield_bounds_size: wp.array(dtype=wp.vec3), hfield_bounds_size: wp.array[wp.vec3],
# Out: # Out:
lower_out: wp.array(dtype=wp.vec3), lower_out: wp.array[wp.vec3],
upper_out: wp.array(dtype=wp.vec3), upper_out: wp.array[wp.vec3],
group_out: wp.array(dtype=int), group_out: wp.array[int],
): ):
worldid, geom_local_id = wp.tid() worldid, geom_local_id = wp.tid()
geom_id = enabled_geom_ids[geom_local_id] geom_id = enabled_geom_ids[geom_local_id]
@@ -205,8 +205,13 @@ def _compute_bvh_bounds(
elif type == GeomType.PLANE: elif type == GeomType.PLANE:
lower_bound, upper_bound = _compute_plane_bounds(pos, rot, size) lower_bound, upper_bound = _compute_plane_bounds(pos, rot, size)
elif type == GeomType.MESH: elif type == GeomType.MESH:
size = mesh_bounds_size[geom_dataid[geom_id]] did = geom_dataid[worldid % geom_dataid.shape[0], geom_id]
if did >= 0:
size = mesh_bounds_size[did]
lower_bound, upper_bound = _compute_box_bounds(pos, rot, size) lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
else:
lower_bound = pos
upper_bound = pos
elif type == GeomType.ELLIPSOID: elif type == GeomType.ELLIPSOID:
lower_bound, upper_bound = _compute_ellipsoid_bounds(pos, rot, size) lower_bound, upper_bound = _compute_ellipsoid_bounds(pos, rot, size)
elif type == GeomType.CYLINDER: elif type == GeomType.CYLINDER:
@@ -214,9 +219,14 @@ def _compute_bvh_bounds(
elif type == GeomType.BOX: elif type == GeomType.BOX:
lower_bound, upper_bound = _compute_box_bounds(pos, rot, size) lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
elif type == GeomType.HFIELD: elif type == GeomType.HFIELD:
size = hfield_bounds_size[geom_dataid[geom_id]] did = geom_dataid[worldid % geom_dataid.shape[0], geom_id]
if did >= 0:
size = hfield_bounds_size[did]
hfield_center = pos + rot[:, 2] * size[2] hfield_center = pos + rot[:, 2] * size[2]
lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size) lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size)
else:
lower_bound = pos
upper_bound = pos
lower_out[worldid * bvh_ngeom + geom_local_id] = lower_bound lower_out[worldid * bvh_ngeom + geom_local_id] = lower_bound
upper_out[worldid * bvh_ngeom + geom_local_id] = upper_bound upper_out[worldid * bvh_ngeom + geom_local_id] = upper_bound
@@ -228,7 +238,7 @@ def compute_bvh_group_roots(
# In: # In:
bvh_id: wp.uint64, bvh_id: wp.uint64,
# Out: # Out:
group_root_out: wp.array(dtype=int), group_root_out: wp.array[int],
): ):
tid = wp.tid() tid = wp.tid()
root = wp.bvh_get_group_root(bvh_id, tid) root = wp.bvh_get_group_root(bvh_id, tid)
@@ -238,21 +248,21 @@ def compute_bvh_group_roots(
@wp.kernel @wp.kernel
def _compute_flex_bvh_bounds( def _compute_flex_bvh_bounds(
# Model: # Model:
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_vertnum: wp.array(dtype=int), flex_vertnum: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
flex_geom_flexid: wp.array(dtype=int), flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array(dtype=int), flex_geom_edgeid: wp.array[int],
bvh_ngeom: int, bvh_ngeom: int,
total_bvh_size: int, total_bvh_size: int,
# Out: # Out:
lower_out: wp.array(dtype=wp.vec3), lower_out: wp.array[wp.vec3],
upper_out: wp.array(dtype=wp.vec3), upper_out: wp.array[wp.vec3],
group_out: wp.array(dtype=int), group_out: wp.array[int],
): ):
worldid, flexlocalid = wp.tid() worldid, flexlocalid = wp.tid()
@@ -289,7 +299,7 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom
geom_type = wp.array(mjm.geom_type, dtype=int) geom_type = wp.array(mjm.geom_type, dtype=int)
geom_dataid = wp.array(mjm.geom_dataid, dtype=int) geom_dataid = wp.array(np.tile(mjm.geom_dataid, (nworld, 1)), dtype=int)
geom_size = wp.array(np.tile(mjm.geom_size[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3) 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_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) geom_xmat = wp.array(np.tile(mjd.geom_xmat.reshape(mjm.ngeom, 3, 3)[np.newaxis, :, :, :], (nworld, 1, 1, 1)), dtype=wp.mat33)
@@ -599,16 +609,16 @@ def build_hfield_bvh(
def accumulate_flex_vertex_normals( def accumulate_flex_vertex_normals(
# Model: # Model:
nflex: int, nflex: int,
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_elemadr: wp.array(dtype=int), flex_elemadr: wp.array[int],
flex_elemnum: wp.array(dtype=int), flex_elemnum: wp.array[int],
flex_elemdataadr: wp.array(dtype=int), flex_elemdataadr: wp.array[int],
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# Out: # Out:
flexvert_norm_out: wp.array2d(dtype=wp.vec3), flexvert_norm_out: wp.array2d[wp.vec3],
): ):
"""Accumulate per-vertex normals by summing adjacent face normals.""" """Accumulate per-vertex normals by summing adjacent face normals."""
worldid, elemid = wp.tid() worldid, elemid = wp.tid()
@@ -644,7 +654,7 @@ def accumulate_flex_vertex_normals(
@wp.kernel @wp.kernel
def normalize_vertex_normals( def normalize_vertex_normals(
# Out: # Out:
flexvert_norm_out: wp.array2d(dtype=wp.vec3), flexvert_norm_out: wp.array2d[wp.vec3],
): ):
"""Normalize accumulated vertex normals.""" """Normalize accumulated vertex normals."""
worldid, vertid = wp.tid() worldid, vertid = wp.tid()
@@ -654,20 +664,20 @@ def normalize_vertex_normals(
@wp.kernel @wp.kernel
def _build_flex_2d_elements( def _build_flex_2d_elements(
# Model: # Model:
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
flexvert_norm_in: wp.array2d(dtype=wp.vec3), flexvert_norm_in: wp.array2d[wp.vec3],
elem_adr: int, elem_adr: int,
vert_adr: int, vert_adr: int,
face_offset: int, face_offset: int,
radius: float, radius: float,
nfaces: int, nfaces: int,
# Out: # Out:
face_point_out: wp.array(dtype=wp.vec3), face_point_out: wp.array[wp.vec3],
face_index_out: wp.array(dtype=int), face_index_out: wp.array[int],
group_out: wp.array(dtype=int), group_out: wp.array[int],
): ):
"""Create faces from 2D flex elements (triangles). """Create faces from 2D flex elements (triangles).
@@ -728,20 +738,20 @@ def _build_flex_2d_elements(
@wp.kernel @wp.kernel
def _build_flex_2d_sides( def _build_flex_2d_sides(
# Model: # Model:
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
flexvert_norm_in: wp.array2d(dtype=wp.vec3), flexvert_norm_in: wp.array2d[wp.vec3],
shell_adr: int, shell_adr: int,
vert_adr: int, vert_adr: int,
face_offset: int, face_offset: int,
radius: float, radius: float,
nface: int, nface: int,
# Out: # Out:
face_point_out: wp.array(dtype=wp.vec3), face_point_out: wp.array[wp.vec3],
face_index_out: wp.array(dtype=int), face_index_out: wp.array[int],
group_out: wp.array(dtype=int), group_out: wp.array[int],
): ):
"""Create side faces from 2D flex shell fragments. """Create side faces from 2D flex shell fragments.
@@ -790,18 +800,18 @@ def _build_flex_2d_sides(
@wp.kernel @wp.kernel
def _build_flex_3d_shells( def _build_flex_3d_shells(
# Model: # Model:
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
shell_adr: int, shell_adr: int,
vert_adr: int, vert_adr: int,
face_offset: int, face_offset: int,
nface: int, nface: int,
# Out: # Out:
face_point_out: wp.array(dtype=wp.vec3), face_point_out: wp.array[wp.vec3],
face_index_out: wp.array(dtype=int), face_index_out: wp.array[int],
group_out: wp.array(dtype=int), group_out: wp.array[int],
): ):
"""Create faces from 3D flex shell fragments (triangles). """Create faces from 3D flex shell fragments (triangles).
@@ -836,22 +846,22 @@ def _build_flex_3d_shells(
@wp.kernel @wp.kernel
def _update_flex_2d_face_points( def _update_flex_2d_face_points(
# Model: # Model:
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_elemnum: wp.array(dtype=int), flex_elemnum: wp.array[int],
flex_elemdataadr: wp.array(dtype=int), flex_elemdataadr: wp.array[int],
flex_shelldataadr: wp.array(dtype=int), flex_shelldataadr: wp.array[int],
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
flexvert_norm_in: wp.array2d(dtype=wp.vec3), flexvert_norm_in: wp.array2d[wp.vec3],
flex_id: int, flex_id: int,
nface: int, nface: int,
smooth: bool, smooth: bool,
# Out: # Out:
face_point_out: wp.array(dtype=wp.vec3), face_point_out: wp.array[wp.vec3],
): ):
worldid, workid = wp.tid() worldid, workid = wp.tid()
@@ -935,16 +945,16 @@ def _update_flex_2d_face_points(
@wp.kernel @wp.kernel
def _update_flex_3d_face_points( def _update_flex_3d_face_points(
# Model: # Model:
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_shelldataadr: wp.array(dtype=int), flex_shelldataadr: wp.array[int],
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
flex_id: int, flex_id: int,
nface: int, nface: int,
# Out: # Out:
face_point_out: wp.array(dtype=wp.vec3), face_point_out: wp.array[wp.vec3],
): ):
worldid, shellid = wp.tid() worldid, shellid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+189 -188
View File
@@ -55,20 +55,20 @@ mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
@wp.func @wp.func
def _hfield_filter( def _hfield_filter(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
hfield_size: wp.array(dtype=wp.vec4), hfield_size: wp.array[wp.vec4],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
worldid: int, worldid: int,
g1: int, g1: int,
@@ -79,7 +79,8 @@ def _hfield_filter(
See MuJoCo mjc_ConvexHField. See MuJoCo mjc_ConvexHField.
""" """
# height field info # height field info
hfdataid = geom_dataid[g1] dataid_setid = worldid % geom_dataid.shape[0]
hfdataid = geom_dataid[dataid_setid, g1]
size1 = hfield_size[hfdataid] size1 = hfield_size[hfdataid]
# geom info # geom info
@@ -124,7 +125,7 @@ def _hfield_filter(
# load mesh vertex data for support function queries # load mesh vertex data for support function queries
if geomtype2 == GeomType.MESH: if geomtype2 == GeomType.MESH:
dataid = geom_dataid[g2] dataid = geom_dataid[dataid_setid, g2]
geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1) geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1) geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1) geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
@@ -167,78 +168,78 @@ def ccd_hfield_kernel_builder(
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def ccd_hfield_kernel( def ccd_hfield_kernel(
# Model: # Model:
opt_ccd_tolerance: wp.array(dtype=float), opt_ccd_tolerance: wp.array[float],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_priority: wp.array(dtype=int), geom_priority: wp.array[int],
geom_solmix: wp.array2d(dtype=float), geom_solmix: wp.array2d[float],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d[float],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
mesh_polynum: wp.array(dtype=int), mesh_polynum: wp.array[int],
mesh_polyadr: wp.array(dtype=int), mesh_polyadr: wp.array[int],
mesh_polynormal: wp.array(dtype=wp.vec3), mesh_polynormal: wp.array[wp.vec3],
mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertadr: wp.array[int],
mesh_polyvertnum: wp.array(dtype=int), mesh_polyvertnum: wp.array[int],
mesh_polyvert: wp.array(dtype=int), mesh_polyvert: wp.array[int],
mesh_polymapadr: wp.array(dtype=int), mesh_polymapadr: wp.array[int],
mesh_polymapnum: wp.array(dtype=int), mesh_polymapnum: wp.array[int],
mesh_polymap: wp.array(dtype=int), mesh_polymap: wp.array[int],
hfield_size: wp.array(dtype=wp.vec4), hfield_size: wp.array[wp.vec4],
hfield_nrow: wp.array(dtype=int), hfield_nrow: wp.array[int],
hfield_ncol: wp.array(dtype=int), hfield_ncol: wp.array[int],
hfield_adr: wp.array(dtype=int), hfield_adr: wp.array[int],
hfield_data: wp.array(dtype=float), hfield_data: wp.array[float],
pair_dim: wp.array(dtype=int), pair_dim: wp.array[int],
pair_solref: wp.array2d(dtype=wp.vec2), pair_solref: wp.array2d[wp.vec2],
pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d[wp.vec2],
pair_solimp: wp.array2d(dtype=vec5), pair_solimp: wp.array2d[vec5],
pair_margin: wp.array2d(dtype=float), pair_margin: wp.array2d[float],
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d[float],
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d[vec5],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
naconmax_in: int, naconmax_in: int,
naccdmax_in: int, naccdmax_in: int,
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array[int],
# In: # In:
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array[wp.vec2i],
collision_pairid_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array[wp.vec2i],
collision_worldid_in: wp.array(dtype=int), collision_worldid_in: wp.array[int],
epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert_in: wp.array2d[wp.vec3],
epa_vert_index_in: wp.array2d(dtype=int), epa_vert_index_in: wp.array2d[int],
epa_face_in: wp.array2d(dtype=int), epa_face_in: wp.array2d[int],
epa_pr_in: wp.array2d(dtype=wp.vec3), epa_pr_in: wp.array2d[wp.vec3],
epa_norm2_in: wp.array2d(dtype=float), epa_norm2_in: wp.array2d[float],
epa_horizon_in: wp.array2d(dtype=int), epa_horizon_in: wp.array2d[int],
nccd_in: wp.array(dtype=int), nccd_in: wp.array[int],
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
collisionid = wp.tid() collisionid = wp.tid()
if collisionid >= ncollision_in[0]: if collisionid >= ncollision_in[0]:
@@ -339,7 +340,7 @@ def ccd_hfield_kernel_builder(
geom1.rot = wp.identity(n=3, dtype=float) geom1.rot = wp.identity(n=3, dtype=float)
# see MuJoCo mjc_ConvexHField # see MuJoCo mjc_ConvexHField
geom1_dataid = geom_dataid[g1] geom1_dataid = geom_dataid[worldid % geom_dataid.shape[0], g1]
# height field subgrid # height field subgrid
nrow = hfield_nrow[geom1_dataid] nrow = hfield_nrow[geom1_dataid]
@@ -710,27 +711,27 @@ def ccd_kernel_builder(
@wp.func @wp.func
def eval_ccd_write_contact( def eval_ccd_write_contact(
# Model: # Model:
opt_ccd_tolerance: wp.array(dtype=float), opt_ccd_tolerance: wp.array[float],
# Data in: # Data in:
naconmax_in: int, naconmax_in: int,
# In: # In:
epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert_in: wp.array2d[wp.vec3],
epa_vert_index_in: wp.array2d(dtype=int), epa_vert_index_in: wp.array2d[int],
epa_face_in: wp.array2d(dtype=int), epa_face_in: wp.array2d[int],
epa_pr_in: wp.array2d(dtype=wp.vec3), epa_pr_in: wp.array2d[wp.vec3],
epa_norm2_in: wp.array2d(dtype=float), epa_norm2_in: wp.array2d[float],
epa_horizon_in: wp.array2d(dtype=int), epa_horizon_in: wp.array2d[int],
multiccd_polygon_in: wp.array2d(dtype=wp.vec3), multiccd_polygon_in: wp.array2d[wp.vec3],
multiccd_clipped_in: wp.array2d(dtype=wp.vec3), multiccd_clipped_in: wp.array2d[wp.vec3],
multiccd_pnormal_in: wp.array2d(dtype=wp.vec3), multiccd_pnormal_in: wp.array2d[wp.vec3],
multiccd_pdist_in: wp.array2d(dtype=float), multiccd_pdist_in: wp.array2d[float],
multiccd_idx1_in: wp.array2d(dtype=int), multiccd_idx1_in: wp.array2d[int],
multiccd_idx2_in: wp.array2d(dtype=int), multiccd_idx2_in: wp.array2d[int],
multiccd_n1_in: wp.array2d(dtype=wp.vec3), multiccd_n1_in: wp.array2d[wp.vec3],
multiccd_n2_in: wp.array2d(dtype=wp.vec3), multiccd_n2_in: wp.array2d[wp.vec3],
multiccd_endvert_in: wp.array2d(dtype=wp.vec3), multiccd_endvert_in: wp.array2d[wp.vec3],
multiccd_face1_in: wp.array2d(dtype=wp.vec3), multiccd_face1_in: wp.array2d[wp.vec3],
multiccd_face2_in: wp.array2d(dtype=wp.vec3), multiccd_face2_in: wp.array2d[wp.vec3],
geom1: Geom, geom1: Geom,
geom2: Geom, geom2: Geom,
geoms: wp.vec2i, geoms: wp.vec2i,
@@ -747,21 +748,21 @@ def ccd_kernel_builder(
x2: wp.vec3, x2: wp.vec3,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
) -> int: ) -> int:
points = mat43() points = mat43()
witness1 = mat43() witness1 = mat43()
@@ -892,83 +893,83 @@ def ccd_kernel_builder(
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def ccd_kernel( def ccd_kernel(
# Model: # Model:
opt_ccd_tolerance: wp.array(dtype=float), opt_ccd_tolerance: wp.array[float],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_priority: wp.array(dtype=int), geom_priority: wp.array[int],
geom_solmix: wp.array2d(dtype=float), geom_solmix: wp.array2d[float],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d[float],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
mesh_polynum: wp.array(dtype=int), mesh_polynum: wp.array[int],
mesh_polyadr: wp.array(dtype=int), mesh_polyadr: wp.array[int],
mesh_polynormal: wp.array(dtype=wp.vec3), mesh_polynormal: wp.array[wp.vec3],
mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertadr: wp.array[int],
mesh_polyvertnum: wp.array(dtype=int), mesh_polyvertnum: wp.array[int],
mesh_polyvert: wp.array(dtype=int), mesh_polyvert: wp.array[int],
mesh_polymapadr: wp.array(dtype=int), mesh_polymapadr: wp.array[int],
mesh_polymapnum: wp.array(dtype=int), mesh_polymapnum: wp.array[int],
mesh_polymap: wp.array(dtype=int), mesh_polymap: wp.array[int],
pair_dim: wp.array(dtype=int), pair_dim: wp.array[int],
pair_solref: wp.array2d(dtype=wp.vec2), pair_solref: wp.array2d[wp.vec2],
pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d[wp.vec2],
pair_solimp: wp.array2d(dtype=vec5), pair_solimp: wp.array2d[vec5],
pair_margin: wp.array2d(dtype=float), pair_margin: wp.array2d[float],
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d[float],
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d[vec5],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
naconmax_in: int, naconmax_in: int,
naccdmax_in: int, naccdmax_in: int,
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array[int],
# In: # In:
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array[wp.vec2i],
collision_pairid_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array[wp.vec2i],
collision_worldid_in: wp.array(dtype=int), collision_worldid_in: wp.array[int],
epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert_in: wp.array2d[wp.vec3],
epa_vert_index_in: wp.array2d(dtype=int), epa_vert_index_in: wp.array2d[int],
epa_face_in: wp.array2d(dtype=int), epa_face_in: wp.array2d[int],
epa_pr_in: wp.array2d(dtype=wp.vec3), epa_pr_in: wp.array2d[wp.vec3],
epa_norm2_in: wp.array2d(dtype=float), epa_norm2_in: wp.array2d[float],
epa_horizon_in: wp.array2d(dtype=int), epa_horizon_in: wp.array2d[int],
multiccd_polygon_in: wp.array2d(dtype=wp.vec3), multiccd_polygon_in: wp.array2d[wp.vec3],
multiccd_clipped_in: wp.array2d(dtype=wp.vec3), multiccd_clipped_in: wp.array2d[wp.vec3],
multiccd_pnormal_in: wp.array2d(dtype=wp.vec3), multiccd_pnormal_in: wp.array2d[wp.vec3],
multiccd_pdist_in: wp.array2d(dtype=float), multiccd_pdist_in: wp.array2d[float],
multiccd_idx1_in: wp.array2d(dtype=int), multiccd_idx1_in: wp.array2d[int],
multiccd_idx2_in: wp.array2d(dtype=int), multiccd_idx2_in: wp.array2d[int],
multiccd_n1_in: wp.array2d(dtype=wp.vec3), multiccd_n1_in: wp.array2d[wp.vec3],
multiccd_n2_in: wp.array2d(dtype=wp.vec3), multiccd_n2_in: wp.array2d[wp.vec3],
multiccd_endvert_in: wp.array2d(dtype=wp.vec3), multiccd_endvert_in: wp.array2d[wp.vec3],
multiccd_face1_in: wp.array2d(dtype=wp.vec3), multiccd_face1_in: wp.array2d[wp.vec3],
multiccd_face2_in: wp.array2d(dtype=wp.vec3), multiccd_face2_in: wp.array2d[wp.vec3],
nccd_in: wp.array(dtype=int), nccd_in: wp.array[int],
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
collisionid = wp.tid() collisionid = wp.tid()
if collisionid >= ncollision_in[0]: if collisionid >= ncollision_in[0]:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+64 -62
View File
@@ -47,44 +47,44 @@ class Geom:
hfprism: mat63 hfprism: mat63
vertadr: int vertadr: int
vertnum: int vertnum: int
vert: wp.array(dtype=wp.vec3) vert: wp.array[wp.vec3]
graphadr: int graphadr: int
graph: wp.array(dtype=int) graph: wp.array[int]
mesh_polynum: int mesh_polynum: int
mesh_polyadr: int mesh_polyadr: int
mesh_polynormal: wp.array(dtype=wp.vec3) mesh_polynormal: wp.array[wp.vec3]
mesh_polyvertadr: wp.array(dtype=int) mesh_polyvertadr: wp.array[int]
mesh_polyvertnum: wp.array(dtype=int) mesh_polyvertnum: wp.array[int]
mesh_polyvert: wp.array(dtype=int) mesh_polyvert: wp.array[int]
mesh_polymapadr: wp.array(dtype=int) mesh_polymapadr: wp.array[int]
mesh_polymapnum: wp.array(dtype=int) mesh_polymapnum: wp.array[int]
mesh_polymap: wp.array(dtype=int) mesh_polymap: wp.array[int]
index: int index: int
@wp.func @wp.func
def geom_collision_pair( def geom_collision_pair(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
mesh_polynum: wp.array(dtype=int), mesh_polynum: wp.array[int],
mesh_polyadr: wp.array(dtype=int), mesh_polyadr: wp.array[int],
mesh_polynormal: wp.array(dtype=wp.vec3), mesh_polynormal: wp.array[wp.vec3],
mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertadr: wp.array[int],
mesh_polyvertnum: wp.array(dtype=int), mesh_polyvertnum: wp.array[int],
mesh_polyvert: wp.array(dtype=int), mesh_polyvert: wp.array[int],
mesh_polymapadr: wp.array(dtype=int), mesh_polymapadr: wp.array[int],
mesh_polymapnum: wp.array(dtype=int), mesh_polymapnum: wp.array[int],
mesh_polymap: wp.array(dtype=int), mesh_polymap: wp.array[int],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
geoms: wp.vec2i, geoms: wp.vec2i,
worldid: int, worldid: int,
@@ -109,8 +109,10 @@ def geom_collision_pair(
# z-axis of the rotation matrix, used as the surface normal for plane collisions # z-axis of the rotation matrix, used as the surface normal for plane collisions
geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2]) geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2])
dataid_setid = worldid % geom_dataid.shape[0]
if geom_type1 == GeomType.MESH: if geom_type1 == GeomType.MESH:
dataid = geom_dataid[g1] dataid = geom_dataid[dataid_setid, g1]
geom1.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1) geom1.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
geom1.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1) geom1.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
geom1.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1) geom1.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
@@ -128,7 +130,7 @@ def geom_collision_pair(
geom1.mesh_polymap = mesh_polymap geom1.mesh_polymap = mesh_polymap
if geom_type2 == GeomType.MESH: if geom_type2 == GeomType.MESH:
dataid = geom_dataid[g2] dataid = geom_dataid[dataid_setid, g2]
geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1) geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1) geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1) geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
@@ -174,21 +176,21 @@ def write_contact(
pairid_in: wp.vec2i, pairid_in: wp.vec2i,
worldid_in: int, worldid_in: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
) -> int: ) -> int:
"""Atomically write a detected contact into the contact output arrays. """Atomically write a detected contact into the contact output arrays.
@@ -233,24 +235,24 @@ def write_contact(
@wp.func @wp.func
def contact_params( def contact_params(
# Model: # Model:
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_priority: wp.array(dtype=int), geom_priority: wp.array[int],
geom_solmix: wp.array2d(dtype=float), geom_solmix: wp.array2d[float],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d[float],
pair_dim: wp.array(dtype=int), pair_dim: wp.array[int],
pair_solref: wp.array2d(dtype=wp.vec2), pair_solref: wp.array2d[wp.vec2],
pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d[wp.vec2],
pair_solimp: wp.array2d(dtype=vec5), pair_solimp: wp.array2d[vec5],
pair_margin: wp.array2d(dtype=float), pair_margin: wp.array2d[float],
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d[float],
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d[vec5],
# In: # In:
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array[wp.vec2i],
collision_pairid_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array[wp.vec2i],
cid: int, cid: int,
worldid: int, worldid: int,
): ):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+54 -54
View File
@@ -80,8 +80,8 @@ MJ_COLLISION_TABLE = {
@wp.kernel @wp.kernel
def _zero_nacon_ncollision( def _zero_nacon_ncollision(
# Data out: # Data out:
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
ncollision_out: wp.array(dtype=int), ncollision_out: wp.array[int],
): ):
ncollision_out[0] = 0 ncollision_out[0] = 0
nacon_out[0] = 0 nacon_out[0] = 0
@@ -275,12 +275,12 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
@wp.func @wp.func
def func( def func(
# Model: # Model:
geom_aabb: wp.array3d(dtype=wp.vec3), geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
geom1: int, geom1: int,
geom2: int, geom2: int,
@@ -324,8 +324,8 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
@wp.func @wp.func
def _add_geom_pair( def _add_geom_pair(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
nxn_pairid: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array[wp.vec2i],
# Data in: # Data in:
naconmax_in: int, naconmax_in: int,
# In: # In:
@@ -334,11 +334,11 @@ def _add_geom_pair(
worldid: int, worldid: int,
nxnid: int, nxnid: int,
# Data out: # Data out:
ncollision_out: wp.array(dtype=int), ncollision_out: wp.array[int],
# Out: # Out:
collision_pair_out: wp.array(dtype=wp.vec2i), collision_pair_out: wp.array[wp.vec2i],
collision_pairid_out: wp.array(dtype=wp.vec2i), collision_pairid_out: wp.array[wp.vec2i],
collision_worldid_out: wp.array(dtype=int), collision_worldid_out: wp.array[int],
): ):
pairid = wp.atomic_add(ncollision_out, 0, 1) pairid = wp.atomic_add(ncollision_out, 0, 1)
@@ -359,7 +359,7 @@ def _add_geom_pair(
@wp.func @wp.func
def _binary_search(values: wp.array(dtype=Any), value: Any, lower: int, upper: int) -> int: def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) -> int:
while lower < upper: while lower < upper:
mid = (lower + upper) >> 1 mid = (lower + upper) >> 1
if values[mid] > value: if values[mid] > value:
@@ -375,18 +375,18 @@ def _sap_project(opt_broadphase: int):
def sap_project( def sap_project(
# Model: # Model:
ngeom: int, ngeom: int,
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
nworld_in: int, nworld_in: int,
# In: # In:
direction_in: wp.vec3, direction_in: wp.vec3,
# Out: # Out:
projection_lower_out: wp.array2d(dtype=float), projection_lower_out: wp.array2d[float],
projection_upper_out: wp.array2d(dtype=float), projection_upper_out: wp.array2d[float],
sort_index_out: wp.array2d(dtype=int), sort_index_out: wp.array2d[int],
segmented_index_out: wp.array(dtype=int), segmented_index_out: wp.array[int],
): ):
worldid, geomid = wp.tid() worldid, geomid = wp.tid()
@@ -422,11 +422,11 @@ def _sap_range(
# Model: # Model:
ngeom: int, ngeom: int,
# In: # In:
projection_lower_in: wp.array2d(dtype=float), projection_lower_in: wp.array2d[float],
projection_upper_in: wp.array2d(dtype=float), projection_upper_in: wp.array2d[float],
sort_index_in: wp.array2d(dtype=int), sort_index_in: wp.array2d[int],
# Out: # Out:
range_out: wp.array2d(dtype=int), range_out: wp.array2d[int],
): ):
worldid, geomid = wp.tid() worldid, geomid = wp.tid()
@@ -448,26 +448,26 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
def kernel( def kernel(
# Model: # Model:
ngeom: int, ngeom: int,
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_aabb: wp.array3d(dtype=wp.vec3), geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
nxn_pairid: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array[wp.vec2i],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
nworld_in: int, nworld_in: int,
naconmax_in: int, naconmax_in: int,
# In: # In:
sort_index_in: wp.array2d(dtype=int), sort_index_in: wp.array2d[int],
cumulative_sum_in: wp.array(dtype=int), cumulative_sum_in: wp.array[int],
nsweep_in: int, nsweep_in: int,
# Data out: # Data out:
ncollision_out: wp.array(dtype=int), ncollision_out: wp.array[int],
# Out: # Out:
collision_pair_out: wp.array(dtype=wp.vec2i), collision_pair_out: wp.array[wp.vec2i],
collision_pairid_out: wp.array(dtype=wp.vec2i), collision_pairid_out: wp.array[wp.vec2i],
collision_worldid_out: wp.array(dtype=int), collision_worldid_out: wp.array[int],
): ):
worldgeomid = wp.tid() worldgeomid = wp.tid()
@@ -528,11 +528,11 @@ def _segmented_sort(tile_size: int):
@wp.kernel(module="unique") @wp.kernel(module="unique")
def segmented_sort( def segmented_sort(
# In: # In:
projection_lower_in: wp.array2d(dtype=float), projection_lower_in: wp.array2d[float],
sort_index_in: wp.array2d(dtype=int), sort_index_in: wp.array2d[int],
# Out: # Out:
projection_lower_out: wp.array2d(dtype=float), projection_lower_out: wp.array2d[float],
sort_index_out: wp.array2d(dtype=int), sort_index_out: wp.array2d[int],
): ):
worldid = wp.tid() worldid = wp.tid()
@@ -648,22 +648,22 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def kernel( def kernel(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_aabb: wp.array3d(dtype=wp.vec3), geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
nxn_geom_pair: wp.array(dtype=wp.vec2i), nxn_geom_pair: wp.array[wp.vec2i],
nxn_pairid: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array[wp.vec2i],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
naconmax_in: int, naconmax_in: int,
# Data out: # Data out:
ncollision_out: wp.array(dtype=int), ncollision_out: wp.array[int],
# Out: # Out:
collision_pair_out: wp.array(dtype=wp.vec2i), collision_pair_out: wp.array[wp.vec2i],
collision_pairid_out: wp.array(dtype=wp.vec2i), collision_pairid_out: wp.array[wp.vec2i],
collision_worldid_out: wp.array(dtype=int), collision_worldid_out: wp.array[int],
): ):
worldid, elementid = wp.tid() worldid, elementid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py
Vendored
+138 -138
View File
@@ -46,22 +46,22 @@ def _write_flex_contact(
vertid: int, vertid: int,
worldid: int, worldid: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_flex_out: wp.array(dtype=wp.vec2i), contact_flex_out: wp.array[wp.vec2i],
contact_vert_out: wp.array(dtype=wp.vec2i), contact_vert_out: wp.array[wp.vec2i],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
if dist >= margin or dist >= MJ_MAXVAL: if dist >= margin or dist >= MJ_MAXVAL:
return return
@@ -110,22 +110,22 @@ def _collide_geom_triangle(
vertex_id: int, vertex_id: int,
worldid: int, worldid: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_flex_out: wp.array(dtype=wp.vec2i), contact_flex_out: wp.array[wp.vec2i],
contact_vert_out: wp.array(dtype=wp.vec2i), contact_vert_out: wp.array[wp.vec2i],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
if gtype == int(GeomType.SPHERE): if gtype == int(GeomType.SPHERE):
sphere_radius = size_val[0] sphere_radius = size_val[0]
@@ -262,41 +262,41 @@ def _flex_plane_narrowphase(
# Model: # Model:
ngeom: int, ngeom: int,
nflexvert: int, nflexvert: int,
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
flex_condim: wp.array(dtype=int), flex_condim: wp.array[int],
flex_friction: wp.array(dtype=wp.vec3), flex_friction: wp.array[wp.vec3],
flex_margin: wp.array(dtype=float), flex_margin: wp.array[float],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
flex_vertflexid: wp.array(dtype=int), flex_vertflexid: wp.array[int],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
nworld_in: int, nworld_in: int,
naconmax_in: int, naconmax_in: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_flex_out: wp.array(dtype=wp.vec2i), contact_flex_out: wp.array[wp.vec2i],
contact_vert_out: wp.array(dtype=wp.vec2i), contact_vert_out: wp.array[wp.vec2i],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
worldid, vertid = wp.tid() worldid, vertid = wp.tid()
@@ -382,48 +382,48 @@ def _flex_narrowphase_dim2(
# Model: # Model:
ngeom: int, ngeom: int,
nflex: int, nflex: int,
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_contype: wp.array(dtype=int), geom_contype: wp.array[int],
geom_conaffinity: wp.array(dtype=int), geom_conaffinity: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
flex_contype: wp.array(dtype=int), flex_contype: wp.array[int],
flex_conaffinity: wp.array(dtype=int), flex_conaffinity: wp.array[int],
flex_margin: wp.array(dtype=float), flex_margin: wp.array[float],
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_elemadr: wp.array(dtype=int), flex_elemadr: wp.array[int],
flex_elemnum: wp.array(dtype=int), flex_elemnum: wp.array[int],
flex_elemdataadr: wp.array(dtype=int), flex_elemdataadr: wp.array[int],
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
nworld_in: int, nworld_in: int,
naconmax_in: int, naconmax_in: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_flex_out: wp.array(dtype=wp.vec2i), contact_flex_out: wp.array[wp.vec2i],
contact_vert_out: wp.array(dtype=wp.vec2i), contact_vert_out: wp.array[wp.vec2i],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
worldid, elemid = wp.tid() worldid, elemid = wp.tid()
@@ -533,47 +533,47 @@ def _flex_narrowphase_dim3(
# Model: # Model:
ngeom: int, ngeom: int,
nflex: int, nflex: int,
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_contype: wp.array(dtype=int), geom_contype: wp.array[int],
geom_conaffinity: wp.array(dtype=int), geom_conaffinity: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
flex_contype: wp.array(dtype=int), flex_contype: wp.array[int],
flex_conaffinity: wp.array(dtype=int), flex_conaffinity: wp.array[int],
flex_margin: wp.array(dtype=float), flex_margin: wp.array[float],
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_shellnum: wp.array(dtype=int), flex_shellnum: wp.array[int],
flex_shelldataadr: wp.array(dtype=int), flex_shelldataadr: wp.array[int],
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
nworld_in: int, nworld_in: int,
naconmax_in: int, naconmax_in: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_flex_out: wp.array(dtype=wp.vec2i), contact_flex_out: wp.array[wp.vec2i],
contact_vert_out: wp.array(dtype=wp.vec2i), contact_vert_out: wp.array[wp.vec2i],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
worldid, shellid = wp.tid() worldid, shellid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+68 -74
View File
@@ -63,20 +63,20 @@ class Polytope:
status: int status: int
# vertices in polytope (packed geom1 followed by geom2) # vertices in polytope (packed geom1 followed by geom2)
vert: wp.array(dtype=wp.vec3) vert: wp.array[wp.vec3]
vert_index: wp.array(dtype=int) vert_index: wp.array[int]
nvert: int nvert: int
# faces in polytope # faces in polytope
# 10 bits per each vertex index, while the last significant bits are for # 10 bits per each vertex index, while the last significant bits are for
# invalid and deleted face # invalid and deleted face
face: wp.array(dtype=int) face: wp.array[int]
face_pr: wp.array(dtype=wp.vec3) face_pr: wp.array[wp.vec3]
face_norm2: wp.array(dtype=float) face_norm2: wp.array[float]
nface: int nface: int
# edges that make up the horizon when adding new vertices to polytope # edges that make up the horizon when adding new vertices to polytope
horizon: wp.array(dtype=int) horizon: wp.array[int]
nhorizon: int nhorizon: int
@@ -1335,7 +1335,7 @@ def _area4(a: wp.vec3, b: wp.vec3, c: wp.vec3, d: wp.vec3) -> float:
@wp.func @wp.func
def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i: def _polygon_quad(polygon: wp.array[wp.vec3], npolygon: int) -> wp.vec4i:
"""Returns the indices of a quadrilateral of maximum area in a convex polygon (npolygon > 4).""" """Returns the indices of a quadrilateral of maximum area in a convex polygon (npolygon > 4)."""
b = int(1) b = int(1)
c = int(2) c = int(2)
@@ -1376,7 +1376,7 @@ def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i:
# return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary # return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary
@wp.func @wp.func
def _feature_dim( def _feature_dim(
face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3), offset: int face: wp.vec3i, vert_index: wp.array[int], vert: wp.array[wp.vec3], offset: int
) -> Tuple[int, wp.vec3i, wp.mat33]: ) -> Tuple[int, wp.vec3i, wp.mat33]:
v1i = vert_index[2 * face[0] + offset] v1i = vert_index[2 * face[0] + offset]
v2i = vert_index[2 * face[1] + offset] v2i = vert_index[2 * face[1] + offset]
@@ -1401,9 +1401,7 @@ def _feature_dim(
# find two normals that are facing each other within a tolerance, return 1 if found # find two normals that are facing each other within a tolerance, return 1 if found
@wp.func @wp.func
def _aligned_faces( def _aligned_faces(vert1: wp.array[wp.vec3], len1: int, vert2: wp.array[wp.vec3], len2: int) -> Tuple[int, wp.vec2i]:
vert1: wp.array(dtype=wp.vec3), len1: int, vert2: wp.array(dtype=wp.vec3), len2: int
) -> Tuple[int, wp.vec2i]:
res = wp.vec2i() res = wp.vec2i()
for i in range(len1): for i in range(len1):
for j in range(len2): for j in range(len2):
@@ -1417,9 +1415,7 @@ def _aligned_faces(
# find two normals that are perpendicular to each other within a tolerance # find two normals that are perpendicular to each other within a tolerance
# return 1 if found # return 1 if found
@wp.func @wp.func
def _aligned_face_edge( def _aligned_face_edge(edge: wp.array[wp.vec3], nedge: int, face: wp.array[wp.vec3], nface: int) -> Tuple[int, wp.vec2i]:
edge: wp.array(dtype=wp.vec3), nedge: int, face: wp.array(dtype=wp.vec3), nface: int
) -> Tuple[int, wp.vec2i]:
res = wp.vec2i() res = wp.vec2i()
for i in range(nface): for i in range(nface):
for j in range(nedge): for j in range(nedge):
@@ -1432,9 +1428,7 @@ def _aligned_face_edge(
# find up to n <= 2 common integers of two arrays, return n # find up to n <= 2 common integers of two arrays, return n
@wp.func @wp.func
def _intersect1( def _intersect1(a1: wp.array[int], a2: wp.array[int], start1: int, start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]:
a1: wp.array(dtype=int), a2: wp.array(dtype=int), start1: int, start2: int, len1: int, len2: int
) -> Tuple[int, wp.vec2i]:
count = int(0) count = int(0)
res = wp.vec2i() res = wp.vec2i()
for i in range(start1, start1 + len1): for i in range(start1, start1 + len1):
@@ -1448,7 +1442,7 @@ def _intersect1(
@wp.func @wp.func
def _intersect2(a1: wp.vec2i, a2: wp.array(dtype=int), start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]: def _intersect2(a1: wp.vec2i, a2: wp.array[int], start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]:
count = int(0) count = int(0)
res = wp.vec2i() res = wp.vec2i()
for i in range(len1): for i in range(len1):
@@ -1470,13 +1464,13 @@ def _mesh_normals(
mat: wp.mat33, mat: wp.mat33,
vertadr: int, vertadr: int,
polyadr: int, polyadr: int,
polynormal: wp.array(dtype=wp.vec3), polynormal: wp.array[wp.vec3],
polymapadr: wp.array(dtype=int), polymapadr: wp.array[int],
polymapnum: wp.array(dtype=int), polymapnum: wp.array[int],
polymap: wp.array(dtype=int), polymap: wp.array[int],
# Out: # Out:
normals_out: wp.array(dtype=wp.vec3), normals_out: wp.array[wp.vec3],
indices_out: wp.array(dtype=int), indices_out: wp.array[int],
) -> int: ) -> int:
v1 = feature_index[0] v1 = feature_index[0]
v2 = feature_index[1] v2 = feature_index[1]
@@ -1540,19 +1534,19 @@ def _mesh_edge_normals(
pos: wp.vec3, pos: wp.vec3,
vertadr: int, vertadr: int,
polyadr: int, polyadr: int,
vert: wp.array(dtype=wp.vec3), vert: wp.array[wp.vec3],
polyvertadr: wp.array(dtype=int), polyvertadr: wp.array[int],
polyvertnum: wp.array(dtype=int), polyvertnum: wp.array[int],
polyvert: wp.array(dtype=int), polyvert: wp.array[int],
polymapadr: wp.array(dtype=int), polymapadr: wp.array[int],
polymapnum: wp.array(dtype=int), polymapnum: wp.array[int],
polymap: wp.array(dtype=int), polymap: wp.array[int],
v1: wp.vec3, v1: wp.vec3,
v2: wp.vec3, v2: wp.vec3,
v1i: int, v1i: int,
# Out: # Out:
normals_out: wp.array(dtype=wp.vec3), normals_out: wp.array[wp.vec3],
endverts_out: wp.array(dtype=wp.vec3), endverts_out: wp.array[wp.vec3],
) -> int: ) -> int:
# only one edge # only one edge
if dim == 2: if dim == 2:
@@ -1586,8 +1580,8 @@ def _box_normals2(
mat: wp.mat33, mat: wp.mat33,
n: wp.vec3, n: wp.vec3,
# Out: # Out:
normal_out: wp.array(dtype=wp.vec3), normal_out: wp.array[wp.vec3],
index_out: wp.array(dtype=int), index_out: wp.array[int],
) -> int: ) -> int:
# list of box face normals # list of box face normals
face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0) face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0)
@@ -1621,8 +1615,8 @@ def _box_normals(
mat: wp.mat33, mat: wp.mat33,
dir: wp.vec3, dir: wp.vec3,
# Out: # Out:
normal_out: wp.array(dtype=wp.vec3), normal_out: wp.array[wp.vec3],
index_out: wp.array(dtype=int), index_out: wp.array[int],
) -> int: ) -> int:
v1 = feature_index[0] v1 = feature_index[0]
v2 = feature_index[1] v2 = feature_index[1]
@@ -1669,7 +1663,7 @@ def _box_normals(
# c is 1 if edge is diagonal of a box face # c is 1 if edge is diagonal of a box face
# c is 2 if edge is an external edge of box # c is 2 if edge is an external edge of box
if c == 1 or c == 2: if c == 1 or c == 2:
return 2 return c
return _box_normals2(mat, dir, normal_out, index_out) return _box_normals2(mat, dir, normal_out, index_out)
if feature_dim == 1: if feature_dim == 1:
@@ -1698,8 +1692,8 @@ def _box_edge_normals(
v2: wp.vec3, v2: wp.vec3,
v1i: int, v1i: int,
# Out: # Out:
normal_out: wp.array(dtype=wp.vec3), normal_out: wp.array[wp.vec3],
endvert_out: wp.array(dtype=wp.vec3), endvert_out: wp.array[wp.vec3],
) -> int: ) -> int:
if dim == 2: if dim == 2:
endvert_out[0] = v2 endvert_out[0] = v2
@@ -1726,7 +1720,7 @@ def _box_edge_normals(
# recover face of a box from its index # recover face of a box from its index
@wp.func @wp.func
def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int: def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array[wp.vec3]) -> int:
# compute global coordinates of the box face and face normal # compute global coordinates of the box face and face normal
if idx == 0: # right if idx == 0: # right
face_out[0] = mat @ wp.vec3(size[0], size[1], size[2]) + pos face_out[0] = mat @ wp.vec3(size[0], size[1], size[2]) + pos
@@ -1775,13 +1769,13 @@ def _mesh_face(
pos: wp.vec3, pos: wp.vec3,
vertadr: int, vertadr: int,
polyadr: int, polyadr: int,
vert: wp.array(dtype=wp.vec3), vert: wp.array[wp.vec3],
polyvertadr: wp.array(dtype=int), polyvertadr: wp.array[int],
polyvertnum: wp.array(dtype=int), polyvertnum: wp.array[int],
polyvert: wp.array(dtype=int), polyvert: wp.array[int],
idx: int, idx: int,
# Out: # Out:
face_out: wp.array(dtype=wp.vec3), face_out: wp.array[wp.vec3],
) -> int: ) -> int:
adr = polyvertadr[polyadr + idx] adr = polyvertadr[polyadr + idx]
j = int(0) j = int(0)
@@ -1821,17 +1815,17 @@ def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> float:
@wp.func @wp.func
def _polygon_clip( def _polygon_clip(
# In: # In:
plane_normal: wp.array(dtype=wp.vec3), plane_normal: wp.array[wp.vec3],
plane_dist: wp.array(dtype=float), plane_dist: wp.array[float],
face1: wp.array(dtype=wp.vec3), face1: wp.array[wp.vec3],
nface1: int, nface1: int,
face2: wp.array(dtype=wp.vec3), face2: wp.array[wp.vec3],
nface2: int, nface2: int,
n: wp.vec3, n: wp.vec3,
dir: wp.vec3, dir: wp.vec3,
# Out: # Out:
polygon_out: wp.array(dtype=wp.vec3), polygon_out: wp.array[wp.vec3],
clipped_out: wp.array(dtype=wp.vec3), clipped_out: wp.array[wp.vec3],
) -> Tuple[int, mat43, mat43]: ) -> Tuple[int, mat43, mat43]:
witness1 = mat43() witness1 = mat43()
witness2 = mat43() witness2 = mat43()
@@ -1918,13 +1912,13 @@ def _polygon_clip(
@wp.func @wp.func
def _set_edge( def _set_edge(
# In: # In:
vert1: wp.array(dtype=wp.vec3), vert1: wp.array[wp.vec3],
vert2: wp.array(dtype=wp.vec3), vert2: wp.array[wp.vec3],
start: int, start: int,
end: int, end: int,
offset: int, offset: int,
# Out: # Out:
face_out: wp.array(dtype=wp.vec3), face_out: wp.array[wp.vec3],
) -> int: ) -> int:
face_out[0] = vert1[2 * start + offset] face_out[0] = vert1[2 * start + offset]
face_out[1] = vert2[end] face_out[1] = vert2[end]
@@ -1935,19 +1929,19 @@ def _set_edge(
@wp.func @wp.func
def multicontact( def multicontact(
# In: # In:
polygon: wp.array(dtype=wp.vec3), polygon: wp.array[wp.vec3],
clipped: wp.array(dtype=wp.vec3), clipped: wp.array[wp.vec3],
plane_normal: wp.array(dtype=wp.vec3), plane_normal: wp.array[wp.vec3],
plane_dist: wp.array(dtype=float), plane_dist: wp.array[float],
idx1: wp.array(dtype=int), idx1: wp.array[int],
idx2: wp.array(dtype=int), idx2: wp.array[int],
n1: wp.array(dtype=wp.vec3), n1: wp.array[wp.vec3],
n2: wp.array(dtype=wp.vec3), n2: wp.array[wp.vec3],
endvert: wp.array(dtype=wp.vec3), endvert: wp.array[wp.vec3],
face1: wp.array(dtype=wp.vec3), face1: wp.array[wp.vec3],
face2: wp.array(dtype=wp.vec3), face2: wp.array[wp.vec3],
epa_vert: wp.array(dtype=wp.vec3), epa_vert: wp.array[wp.vec3],
epa_vert_index: wp.array(dtype=int), epa_vert_index: wp.array[int],
epa_face: int, epa_face: int,
x1: wp.vec3, x1: wp.vec3,
x2: wp.vec3, x2: wp.vec3,
@@ -2215,12 +2209,12 @@ def ccd(
geomtype2: int, geomtype2: int,
x_1: wp.vec3, x_1: wp.vec3,
x_2: wp.vec3, x_2: wp.vec3,
vert: wp.array(dtype=wp.vec3), vert: wp.array[wp.vec3],
vert_index: wp.array(dtype=int), vert_index: wp.array[int],
face: wp.array(dtype=int), face: wp.array[int],
face_pr: wp.array(dtype=wp.vec3), face_pr: wp.array[wp.vec3],
face_norm2: wp.array(dtype=float), face_norm2: wp.array[float],
horizon: wp.array(dtype=int), horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]: ) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""General convex collision detection via GJK/EPA.""" """General convex collision detection via GJK/EPA."""
full_margin1 = 0.0 full_margin1 = 0.0
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+248 -248
View File
@@ -295,21 +295,21 @@ def plane_sphere_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contact between a sphere and a plane.""" """Calculates contact between a sphere and a plane."""
normal = plane.normal normal = plane.normal
@@ -367,21 +367,21 @@ def sphere_sphere_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contact between two spheres.""" """Calculates contact between two spheres."""
dist, pos, normal = sphere_sphere(sphere1.pos, sphere1.size[0], sphere2.pos, sphere2.size[0]) dist, pos, normal = sphere_sphere(sphere1.pos, sphere1.size[0], sphere2.pos, sphere2.size[0])
@@ -438,21 +438,21 @@ def sphere_capsule_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates one contact between a sphere and a capsule.""" """Calculates one contact between a sphere and a capsule."""
# capsule axis # capsule axis
@@ -512,21 +512,21 @@ def capsule_capsule_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between two capsules.""" """Calculates contacts between two capsules."""
# capsule axes # capsule axes
@@ -598,21 +598,21 @@ def plane_capsule_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between a capsule and a plane.""" """Calculates contacts between a capsule and a plane."""
# capsule axis # capsule axis
@@ -680,21 +680,21 @@ def plane_ellipsoid_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between an ellipsoid and a plane.""" """Calculates contacts between an ellipsoid and a plane."""
dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, ellipsoid.rot, ellipsoid.size) dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, ellipsoid.rot, ellipsoid.size)
@@ -751,21 +751,21 @@ def plane_box_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between a box and a plane.""" """Calculates contacts between a box and a plane."""
dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size) dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size)
@@ -824,21 +824,21 @@ def plane_convex_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between a plane and a convex object.""" """Calculates contacts between a plane and a convex object."""
dist, pos, normal = plane_convex(plane.normal, plane.pos, convex) dist, pos, normal = plane_convex(plane.normal, plane.pos, convex)
@@ -897,21 +897,21 @@ def sphere_cylinder_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between a sphere and a cylinder.""" """Calculates contacts between a sphere and a cylinder."""
# cylinder axis # cylinder axis
@@ -978,21 +978,21 @@ def plane_cylinder_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between a cylinder and a plane.""" """Calculates contacts between a cylinder and a plane."""
# cylinder axis # cylinder axis
@@ -1061,21 +1061,21 @@ def sphere_box_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, box.rot, box.size) dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, box.rot, box.size)
@@ -1131,21 +1131,21 @@ def capsule_box_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between a capsule and a box.""" """Calculates contacts between a capsule and a box."""
# Extract capsule axis # Extract capsule axis
@@ -1216,21 +1216,21 @@ def box_box_wrapper(
geoms: wp.vec2i, geoms: wp.vec2i,
pairid: wp.vec2i, pairid: wp.vec2i,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
"""Calculates contacts between two boxes.""" """Calculates contacts between two boxes."""
# Call the core function to get contact geometry # Call the core function to get contact geometry
@@ -1302,63 +1302,63 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def primitive_narrowphase( def primitive_narrowphase(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_priority: wp.array(dtype=int), geom_priority: wp.array[int],
geom_solmix: wp.array2d(dtype=float), geom_solmix: wp.array2d[float],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d[float],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
mesh_polynum: wp.array(dtype=int), mesh_polynum: wp.array[int],
mesh_polyadr: wp.array(dtype=int), mesh_polyadr: wp.array[int],
mesh_polynormal: wp.array(dtype=wp.vec3), mesh_polynormal: wp.array[wp.vec3],
mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertadr: wp.array[int],
mesh_polyvertnum: wp.array(dtype=int), mesh_polyvertnum: wp.array[int],
mesh_polyvert: wp.array(dtype=int), mesh_polyvert: wp.array[int],
mesh_polymapadr: wp.array(dtype=int), mesh_polymapadr: wp.array[int],
mesh_polymapnum: wp.array(dtype=int), mesh_polymapnum: wp.array[int],
mesh_polymap: wp.array(dtype=int), mesh_polymap: wp.array[int],
pair_dim: wp.array(dtype=int), pair_dim: wp.array[int],
pair_solref: wp.array2d(dtype=wp.vec2), pair_solref: wp.array2d[wp.vec2],
pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d[wp.vec2],
pair_solimp: wp.array2d(dtype=vec5), pair_solimp: wp.array2d[vec5],
pair_margin: wp.array2d(dtype=float), pair_margin: wp.array2d[float],
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d[float],
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d[vec5],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
naconmax_in: int, naconmax_in: int,
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array[int],
# In: # In:
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array[wp.vec2i],
collision_pairid_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array[wp.vec2i],
collision_worldid_in: wp.array(dtype=int), collision_worldid_in: wp.array[int],
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
tid = wp.tid() tid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+101 -81
View File
@@ -54,9 +54,9 @@ class AABB:
class VolumeData: class VolumeData:
center: wp.vec3 center: wp.vec3
half_size: wp.vec3 half_size: wp.vec3
oct_aabb: wp.array2d(dtype=wp.vec3) oct_aabb: wp.array2d[wp.vec3]
oct_child: wp.array(dtype=vec8i) oct_child: wp.array[vec8i]
oct_coeff: wp.array(dtype=vec8) oct_coeff: wp.array[vec8]
root: int = 0 root: int = 0
valid: bool = False valid: bool = False
@@ -64,10 +64,10 @@ class VolumeData:
@wp.struct @wp.struct
class MeshData: class MeshData:
nmeshface: int nmeshface: int
mesh_vertadr: wp.array(dtype=int) mesh_vertadr: wp.array[int]
mesh_vert: wp.array(dtype=wp.vec3) mesh_vert: wp.array[wp.vec3]
mesh_faceadr: wp.array(dtype=int) mesh_faceadr: wp.array[int]
mesh_face: wp.array(dtype=wp.vec3i) mesh_face: wp.array[wp.vec3i]
data_id: int data_id: int
pos: wp.vec3 pos: wp.vec3
mat: wp.mat33 mat: wp.mat33
@@ -80,12 +80,12 @@ class MeshData:
@wp.func @wp.func
def get_sdf_params( def get_sdf_params(
# Model: # Model:
oct_child: wp.array(dtype=vec8i), oct_child: wp.array[vec8i],
oct_aabb: wp.array2d(dtype=wp.vec3), oct_aabb: wp.array2d[wp.vec3],
oct_coeff: wp.array(dtype=vec8), oct_coeff: wp.array[vec8],
mesh_octadr: wp.array(dtype=int), mesh_octadr: wp.array[int],
plugin: wp.array(dtype=int), plugin: wp.array[int],
plugin_attr: wp.array(dtype=vec_pluginattr), plugin_attr: wp.array[vec_pluginattr],
# In: # In:
g_type: int, g_type: int,
g_size: wp.vec3, g_size: wp.vec3,
@@ -252,7 +252,7 @@ def user_sdf_grad(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> wp.vec3:
@wp.func @wp.func
def find_oct( 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[vec8i], oct_aabb: wp.array2d[wp.vec3], p: wp.vec3, grad: bool, root: int
) -> Tuple[int, Tuple[vec8, vec8, vec8]]: ) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
stack = root stack = root
niter = int(100) niter = int(100)
@@ -665,75 +665,75 @@ def gradient_descent(
def _sdf_narrowphase( def _sdf_narrowphase(
# Model: # Model:
nmeshface: int, nmeshface: int,
oct_child: wp.array(dtype=vec8i), oct_child: wp.array[vec8i],
oct_aabb: wp.array2d(dtype=wp.vec3), oct_aabb: wp.array2d[wp.vec3],
oct_coeff: wp.array(dtype=vec8), oct_coeff: wp.array[vec8],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_priority: wp.array(dtype=int), geom_priority: wp.array[int],
geom_solmix: wp.array2d(dtype=float), geom_solmix: wp.array2d[float],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_solimp: wp.array2d(dtype=vec5), geom_solimp: wp.array2d[vec5],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_aabb: wp.array3d(dtype=wp.vec3), geom_aabb: wp.array3d[wp.vec3],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d[float],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
mesh_octadr: wp.array(dtype=int), mesh_octadr: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_face: wp.array(dtype=wp.vec3i), mesh_face: wp.array[wp.vec3i],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
mesh_polynum: wp.array(dtype=int), mesh_polynum: wp.array[int],
mesh_polyadr: wp.array(dtype=int), mesh_polyadr: wp.array[int],
mesh_polynormal: wp.array(dtype=wp.vec3), mesh_polynormal: wp.array[wp.vec3],
mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertadr: wp.array[int],
mesh_polyvertnum: wp.array(dtype=int), mesh_polyvertnum: wp.array[int],
mesh_polyvert: wp.array(dtype=int), mesh_polyvert: wp.array[int],
mesh_polymapadr: wp.array(dtype=int), mesh_polymapadr: wp.array[int],
mesh_polymapnum: wp.array(dtype=int), mesh_polymapnum: wp.array[int],
mesh_polymap: wp.array(dtype=int), mesh_polymap: wp.array[int],
pair_dim: wp.array(dtype=int), pair_dim: wp.array[int],
pair_solref: wp.array2d(dtype=wp.vec2), pair_solref: wp.array2d[wp.vec2],
pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d[wp.vec2],
pair_solimp: wp.array2d(dtype=vec5), pair_solimp: wp.array2d[vec5],
pair_margin: wp.array2d(dtype=float), pair_margin: wp.array2d[float],
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d[float],
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d[vec5],
plugin: wp.array(dtype=int), plugin: wp.array[int],
plugin_attr: wp.array(dtype=vec_pluginattr), plugin_attr: wp.array[vec_pluginattr],
geom_plugin_index: wp.array(dtype=int), geom_plugin_index: wp.array[int],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
naconmax_in: int, naconmax_in: int,
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array[int],
# In: # In:
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array[wp.vec2i],
collision_pairid_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array[wp.vec2i],
collision_worldid_in: wp.array(dtype=int), collision_worldid_in: wp.array[int],
sdf_initpoints: int, sdf_initpoints: int,
sdf_iterations: int, sdf_iterations: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=vec5), contact_friction_out: wp.array[vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=vec5), contact_solimp_out: wp.array[vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
i, contact_tid = wp.tid() i, contact_tid = wp.tid()
if i >= sdf_initpoints: if i >= sdf_initpoints:
@@ -817,12 +817,32 @@ def _sdf_narrowphase(
pos1 = geom1.pos pos1 = geom1.pos
rot1 = geom1.rot rot1 = geom1.rot
dataid_setid = worldid % geom_dataid.shape[0]
attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params( 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[dataid_setid, g1],
) )
attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params( 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[dataid_setid, g2],
) )
mesh_data1.nmeshface = nmeshface mesh_data1.nmeshface = nmeshface
@@ -830,7 +850,7 @@ def _sdf_narrowphase(
mesh_data1.mesh_vert = mesh_vert mesh_data1.mesh_vert = mesh_vert
mesh_data1.mesh_faceadr = mesh_faceadr mesh_data1.mesh_faceadr = mesh_faceadr
mesh_data1.mesh_face = mesh_face mesh_data1.mesh_face = mesh_face
mesh_data1.data_id = geom_dataid[g1] mesh_data1.data_id = geom_dataid[dataid_setid, g1]
mesh_data1.pos = geom1.pos mesh_data1.pos = geom1.pos
mesh_data1.mat = geom1.rot mesh_data1.mat = geom1.rot
mesh_data1.size = geom1.size mesh_data1.size = geom1.size
@@ -843,7 +863,7 @@ def _sdf_narrowphase(
mesh_data2.mesh_vert = mesh_vert mesh_data2.mesh_vert = mesh_vert
mesh_data2.mesh_faceadr = mesh_faceadr mesh_data2.mesh_faceadr = mesh_faceadr
mesh_data2.mesh_face = mesh_face mesh_data2.mesh_face = mesh_face
mesh_data2.data_id = geom_dataid[g2] mesh_data2.data_id = geom_dataid[dataid_setid, g2]
mesh_data2.pos = geom2.pos mesh_data2.pos = geom2.pos
mesh_data2.mat = geom2.rot mesh_data2.mat = geom2.rot
mesh_data2.size = geom2.size mesh_data2.size = geom2.size
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+400 -400
View File
@@ -32,12 +32,12 @@ wp.set_module_options({"enable_backward": False})
@wp.kernel @wp.kernel
def _zero_constraint_counts( def _zero_constraint_counts(
# Data out: # Data out:
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nf_out: wp.array(dtype=int), nf_out: wp.array[int],
nl_out: wp.array(dtype=int), nl_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid = wp.tid() worldid = wp.tid()
@@ -68,14 +68,14 @@ def _efc_row(
type: int, type: int,
id: int, id: int,
# Out: # Out:
type_out: wp.array2d(dtype=int), type_out: wp.array2d[int],
id_out: wp.array2d(dtype=int), id_out: wp.array2d[int],
pos_out: wp.array2d(dtype=float), pos_out: wp.array2d[float],
margin_out: wp.array2d(dtype=float), margin_out: wp.array2d[float],
D_out: wp.array2d(dtype=float), D_out: wp.array2d[float],
vel_out: wp.array2d(dtype=float), vel_out: wp.array2d[float],
aref_out: wp.array2d(dtype=float), aref_out: wp.array2d[float],
frictionloss_out: wp.array2d(dtype=float), frictionloss_out: wp.array2d[float],
): ):
# calculate kbi # calculate kbi
timeconst = solref[0] timeconst = solref[0]
@@ -126,52 +126,52 @@ def _equality_connect(
# Model: # Model:
nv: int, nv: int,
nsite: int, nsite: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array[int],
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
body_invweight0: wp.array2d(dtype=wp.vec2), body_invweight0: wp.array2d[wp.vec2],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
dof_parentid: wp.array(dtype=int), dof_parentid: wp.array[int],
site_bodyid: wp.array(dtype=int), site_bodyid: wp.array[int],
eq_obj1id: wp.array(dtype=int), eq_obj1id: wp.array[int],
eq_obj2id: wp.array(dtype=int), eq_obj2id: wp.array[int],
eq_objtype: wp.array(dtype=int), eq_objtype: wp.array[int],
eq_solref: wp.array2d(dtype=wp.vec2), eq_solref: wp.array2d[wp.vec2],
eq_solimp: wp.array2d(dtype=vec5), eq_solimp: wp.array2d[vec5],
eq_data: wp.array2d(dtype=vec11), eq_data: wp.array2d[vec11],
is_sparse: bool, is_sparse: bool,
eq_connect_adr: wp.array(dtype=int), eq_connect_adr: wp.array[int],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
eq_active_in: wp.array2d(dtype=bool), eq_active_in: wp.array2d[bool],
xpos_in: wp.array2d(dtype=wp.vec3), xpos_in: wp.array2d[wp.vec3],
xmat_in: wp.array2d(dtype=wp.mat33), xmat_in: wp.array2d[wp.mat33],
site_xpos_in: wp.array2d(dtype=wp.vec3), site_xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
"""Calculates constraint rows for connect equality constraints.""" """Calculates constraint rows for connect equality constraints."""
worldid, eqconnectid = wp.tid() worldid, eqconnectid = wp.tid()
@@ -368,42 +368,42 @@ def _equality_connect(
def _equality_joint( def _equality_joint(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
qpos0: wp.array2d(dtype=float), qpos0: wp.array2d[float],
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
dof_invweight0: wp.array2d(dtype=float), dof_invweight0: wp.array2d[float],
eq_obj1id: wp.array(dtype=int), eq_obj1id: wp.array[int],
eq_obj2id: wp.array(dtype=int), eq_obj2id: wp.array[int],
eq_solref: wp.array2d(dtype=wp.vec2), eq_solref: wp.array2d[wp.vec2],
eq_solimp: wp.array2d(dtype=vec5), eq_solimp: wp.array2d[vec5],
eq_data: wp.array2d(dtype=vec11), eq_data: wp.array2d[vec11],
is_sparse: bool, is_sparse: bool,
eq_jnt_adr: wp.array(dtype=int), eq_jnt_adr: wp.array[int],
# Data in: # Data in:
qpos_in: wp.array2d(dtype=float), qpos_in: wp.array2d[float],
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
eq_active_in: wp.array2d(dtype=bool), eq_active_in: wp.array2d[bool],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, eqjntid = wp.tid() worldid, eqjntid = wp.tid()
eqid = eq_jnt_adr[eqjntid] eqid = eq_jnt_adr[eqjntid]
@@ -499,44 +499,44 @@ def _equality_joint(
def _equality_tendon( def _equality_tendon(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
eq_obj1id: wp.array(dtype=int), eq_obj1id: wp.array[int],
eq_obj2id: wp.array(dtype=int), eq_obj2id: wp.array[int],
eq_solref: wp.array2d(dtype=wp.vec2), eq_solref: wp.array2d[wp.vec2],
eq_solimp: wp.array2d(dtype=vec5), eq_solimp: wp.array2d[vec5],
eq_data: wp.array2d(dtype=vec11), eq_data: wp.array2d[vec11],
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
tendon_length0: wp.array2d(dtype=float), tendon_length0: wp.array2d[float],
tendon_invweight0: wp.array2d(dtype=float), tendon_invweight0: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
eq_ten_adr: wp.array(dtype=int), eq_ten_adr: wp.array[int],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
eq_active_in: wp.array2d(dtype=bool), eq_active_in: wp.array2d[bool],
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
ten_length_in: wp.array2d(dtype=float), ten_length_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, eqtenid = wp.tid() worldid, eqtenid = wp.tid()
eqid = eq_ten_adr[eqtenid] eqid = eq_ten_adr[eqtenid]
@@ -679,42 +679,42 @@ def _equality_flex(is_sparse: bool):
def kernel( def kernel(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
flex_edgeadr: wp.array(dtype=int), flex_edgeadr: wp.array[int],
flex_edgenum: wp.array(dtype=int), flex_edgenum: wp.array[int],
flexedge_length0: wp.array(dtype=float), flexedge_length0: wp.array[float],
flexedge_invweight0: wp.array(dtype=float), flexedge_invweight0: wp.array[float],
flexedge_J_rownnz: wp.array(dtype=int), flexedge_J_rownnz: wp.array[int],
flexedge_J_rowadr: wp.array(dtype=int), flexedge_J_rowadr: wp.array[int],
flexedge_J_colind: wp.array(dtype=int), flexedge_J_colind: wp.array[int],
eq_obj1id: wp.array(dtype=int), eq_obj1id: wp.array[int],
eq_solref: wp.array2d(dtype=wp.vec2), eq_solref: wp.array2d[wp.vec2],
eq_solimp: wp.array2d(dtype=vec5), eq_solimp: wp.array2d[vec5],
eq_flex_adr: wp.array(dtype=int), eq_flex_adr: wp.array[int],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
flexedge_J_in: wp.array2d(dtype=float), flexedge_J_in: wp.array2d[float],
flexedge_length_in: wp.array2d(dtype=float), flexedge_length_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, eqflexid, edgeid = wp.tid() worldid, eqflexid, edgeid = wp.tid()
eqid = eq_flex_adr[eqflexid] eqid = eq_flex_adr[eqflexid]
@@ -794,54 +794,54 @@ def _equality_weld(
# Model: # Model:
nv: int, nv: int,
nsite: int, nsite: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array[int],
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
body_invweight0: wp.array2d(dtype=wp.vec2), body_invweight0: wp.array2d[wp.vec2],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
dof_parentid: wp.array(dtype=int), dof_parentid: wp.array[int],
site_bodyid: wp.array(dtype=int), site_bodyid: wp.array[int],
site_quat: wp.array2d(dtype=wp.quat), site_quat: wp.array2d[wp.quat],
eq_obj1id: wp.array(dtype=int), eq_obj1id: wp.array[int],
eq_obj2id: wp.array(dtype=int), eq_obj2id: wp.array[int],
eq_objtype: wp.array(dtype=int), eq_objtype: wp.array[int],
eq_solref: wp.array2d(dtype=wp.vec2), eq_solref: wp.array2d[wp.vec2],
eq_solimp: wp.array2d(dtype=vec5), eq_solimp: wp.array2d[vec5],
eq_data: wp.array2d(dtype=vec11), eq_data: wp.array2d[vec11],
is_sparse: bool, is_sparse: bool,
eq_wld_adr: wp.array(dtype=int), eq_wld_adr: wp.array[int],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
eq_active_in: wp.array2d(dtype=bool), eq_active_in: wp.array2d[bool],
xpos_in: wp.array2d(dtype=wp.vec3), xpos_in: wp.array2d[wp.vec3],
xquat_in: wp.array2d(dtype=wp.quat), xquat_in: wp.array2d[wp.quat],
xmat_in: wp.array2d(dtype=wp.mat33), xmat_in: wp.array2d[wp.mat33],
site_xpos_in: wp.array2d(dtype=wp.vec3), site_xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, eqweldid = wp.tid() worldid, eqweldid = wp.tid()
eqid = eq_wld_adr[eqweldid] eqid = eq_wld_adr[eqweldid]
@@ -1114,34 +1114,34 @@ def _equality_weld(
def _friction_dof( def _friction_dof(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
dof_solref: wp.array2d(dtype=wp.vec2), dof_solref: wp.array2d[wp.vec2],
dof_solimp: wp.array2d(dtype=vec5), dof_solimp: wp.array2d[vec5],
dof_frictionloss: wp.array2d(dtype=float), dof_frictionloss: wp.array2d[float],
dof_invweight0: wp.array2d(dtype=float), dof_invweight0: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
nf_out: wp.array(dtype=int), nf_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
@@ -1204,38 +1204,38 @@ def _friction_dof(
def _friction_tendon( def _friction_tendon(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
tendon_solref_fri: wp.array2d(dtype=wp.vec2), tendon_solref_fri: wp.array2d[wp.vec2],
tendon_solimp_fri: wp.array2d(dtype=vec5), tendon_solimp_fri: wp.array2d[vec5],
tendon_frictionloss: wp.array2d(dtype=float), tendon_frictionloss: wp.array2d[float],
tendon_invweight0: wp.array2d(dtype=float), tendon_invweight0: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
nf_out: wp.array(dtype=int), nf_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, tenid = wp.tid() worldid, tenid = wp.tid()
@@ -1317,39 +1317,39 @@ def _friction_tendon(
def _limit_slide_hinge( def _limit_slide_hinge(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solref: wp.array2d[wp.vec2],
jnt_solimp: wp.array2d(dtype=vec5), jnt_solimp: wp.array2d[vec5],
jnt_range: wp.array2d(dtype=wp.vec2), jnt_range: wp.array2d[wp.vec2],
jnt_margin: wp.array2d(dtype=float), jnt_margin: wp.array2d[float],
dof_invweight0: wp.array2d(dtype=float), dof_invweight0: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
jnt_limited_slide_hinge_adr: wp.array(dtype=int), jnt_limited_slide_hinge_adr: wp.array[int],
# Data in: # Data in:
qpos_in: wp.array2d(dtype=float), qpos_in: wp.array2d[float],
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
nl_out: wp.array(dtype=int), nl_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, jntlimitedid = wp.tid() worldid, jntlimitedid = wp.tid()
jntid = jnt_limited_slide_hinge_adr[jntlimitedid] jntid = jnt_limited_slide_hinge_adr[jntlimitedid]
@@ -1422,39 +1422,39 @@ def _limit_slide_hinge(
def _limit_ball( def _limit_ball(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solref: wp.array2d[wp.vec2],
jnt_solimp: wp.array2d(dtype=vec5), jnt_solimp: wp.array2d[vec5],
jnt_range: wp.array2d(dtype=wp.vec2), jnt_range: wp.array2d[wp.vec2],
jnt_margin: wp.array2d(dtype=float), jnt_margin: wp.array2d[float],
dof_invweight0: wp.array2d(dtype=float), dof_invweight0: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
jnt_limited_ball_adr: wp.array(dtype=int), jnt_limited_ball_adr: wp.array[int],
# Data in: # Data in:
qpos_in: wp.array2d(dtype=float), qpos_in: wp.array2d[float],
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
nl_out: wp.array(dtype=int), nl_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, jntlimitedid = wp.tid() worldid, jntlimitedid = wp.tid()
jntid = jnt_limited_ball_adr[jntlimitedid] jntid = jnt_limited_ball_adr[jntlimitedid]
@@ -1547,41 +1547,41 @@ def _limit_ball(
def _limit_tendon( def _limit_tendon(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
tendon_solref_lim: wp.array2d(dtype=wp.vec2), tendon_solref_lim: wp.array2d[wp.vec2],
tendon_solimp_lim: wp.array2d(dtype=vec5), tendon_solimp_lim: wp.array2d[vec5],
tendon_range: wp.array2d(dtype=wp.vec2), tendon_range: wp.array2d[wp.vec2],
tendon_margin: wp.array2d(dtype=float), tendon_margin: wp.array2d[float],
tendon_invweight0: wp.array2d(dtype=float), tendon_invweight0: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
tendon_limited_adr: wp.array(dtype=int), tendon_limited_adr: wp.array[int],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
ten_length_in: wp.array2d(dtype=float), ten_length_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
# Data out: # Data out:
nl_out: wp.array(dtype=int), nl_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
worldid, tenlimitedid = wp.tid() worldid, tenlimitedid = wp.tid()
tenid = tendon_limited_adr[tenlimitedid] tenid = tendon_limited_adr[tenlimitedid]
@@ -1669,59 +1669,59 @@ def _limit_tendon(
def _contact_pyramidal( def _contact_pyramidal(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
opt_impratio_invsqrt: wp.array(dtype=float), opt_impratio_invsqrt: wp.array[float],
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array[int],
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
body_invweight0: wp.array2d(dtype=wp.vec2), body_invweight0: wp.array2d[wp.vec2],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
dof_parentid: wp.array(dtype=int), dof_parentid: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_vertbodyid: wp.array(dtype=int), flex_vertbodyid: wp.array[int],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
nacon_in: wp.array(dtype=int), nacon_in: wp.array[int],
# In: # In:
dist_in: wp.array(dtype=float), dist_in: wp.array[float],
condim_in: wp.array(dtype=int), condim_in: wp.array[int],
includemargin_in: wp.array(dtype=float), includemargin_in: wp.array[float],
worldid_in: wp.array(dtype=int), worldid_in: wp.array[int],
geom_in: wp.array(dtype=wp.vec2i), geom_in: wp.array[wp.vec2i],
flex_in: wp.array(dtype=wp.vec2i), flex_in: wp.array[wp.vec2i],
vert_in: wp.array(dtype=wp.vec2i), vert_in: wp.array[wp.vec2i],
pos_in: wp.array(dtype=wp.vec3), pos_in: wp.array[wp.vec3],
frame_in: wp.array(dtype=wp.mat33), frame_in: wp.array[wp.mat33],
friction_in: wp.array(dtype=vec5), friction_in: wp.array[vec5],
solref_in: wp.array(dtype=wp.vec2), solref_in: wp.array[wp.vec2],
solimp_in: wp.array(dtype=vec5), solimp_in: wp.array[vec5],
type_in: wp.array(dtype=int), type_in: wp.array[int],
# Data out: # Data out:
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
conid, dimid = wp.tid() conid, dimid = wp.tid()
@@ -1940,60 +1940,60 @@ def _contact_pyramidal(
def _contact_elliptic( def _contact_elliptic(
# Model: # Model:
nv: int, nv: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
opt_impratio_invsqrt: wp.array(dtype=float), opt_impratio_invsqrt: wp.array[float],
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array[int],
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
body_invweight0: wp.array2d(dtype=wp.vec2), body_invweight0: wp.array2d[wp.vec2],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
dof_parentid: wp.array(dtype=int), dof_parentid: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_vertbodyid: wp.array(dtype=int), flex_vertbodyid: wp.array[int],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
nacon_in: wp.array(dtype=int), nacon_in: wp.array[int],
# In: # In:
dist_in: wp.array(dtype=float), dist_in: wp.array[float],
condim_in: wp.array(dtype=int), condim_in: wp.array[int],
includemargin_in: wp.array(dtype=float), includemargin_in: wp.array[float],
worldid_in: wp.array(dtype=int), worldid_in: wp.array[int],
geom_in: wp.array(dtype=wp.vec2i), geom_in: wp.array[wp.vec2i],
flex_in: wp.array(dtype=wp.vec2i), flex_in: wp.array[wp.vec2i],
vert_in: wp.array(dtype=wp.vec2i), vert_in: wp.array[wp.vec2i],
pos_in: wp.array(dtype=wp.vec3), pos_in: wp.array[wp.vec3],
frame_in: wp.array(dtype=wp.mat33), frame_in: wp.array[wp.mat33],
friction_in: wp.array(dtype=vec5), friction_in: wp.array[vec5],
solref_in: wp.array(dtype=wp.vec2), solref_in: wp.array[wp.vec2],
solreffriction_in: wp.array(dtype=wp.vec2), solreffriction_in: wp.array[wp.vec2],
solimp_in: wp.array(dtype=vec5), solimp_in: wp.array[vec5],
type_in: wp.array(dtype=int), type_in: wp.array[int],
# Data out: # Data out:
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
efc_type_out: wp.array2d(dtype=int), efc_type_out: wp.array2d[int],
efc_id_out: wp.array2d(dtype=int), efc_id_out: wp.array2d[int],
efc_J_rownnz_out: wp.array2d(dtype=int), efc_J_rownnz_out: wp.array2d[int],
efc_J_rowadr_out: wp.array2d(dtype=int), efc_J_rowadr_out: wp.array2d[int],
efc_J_colind_out: wp.array3d(dtype=int), efc_J_colind_out: wp.array3d[int],
efc_J_out: wp.array3d(dtype=float), efc_J_out: wp.array3d[float],
efc_pos_out: wp.array2d(dtype=float), efc_pos_out: wp.array2d[float],
efc_margin_out: wp.array2d(dtype=float), efc_margin_out: wp.array2d[float],
efc_D_out: wp.array2d(dtype=float), efc_D_out: wp.array2d[float],
efc_vel_out: wp.array2d(dtype=float), efc_vel_out: wp.array2d[float],
efc_aref_out: wp.array2d(dtype=float), efc_aref_out: wp.array2d[float],
efc_frictionloss_out: wp.array2d(dtype=float), efc_frictionloss_out: wp.array2d[float],
# Out: # Out:
efc_nnz_out: wp.array(dtype=int), efc_nnz_out: wp.array[int],
): ):
conid, dimid = wp.tid() conid, dimid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+53 -53
View File
@@ -31,27 +31,27 @@ wp.set_module_options({"enable_backward": False})
@wp.kernel @wp.kernel
def _qderiv_actuator_passive_vel( def _qderiv_actuator_passive_vel(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
actuator_dyntype: wp.array(dtype=int), actuator_dyntype: wp.array[int],
actuator_gaintype: wp.array(dtype=int), actuator_gaintype: wp.array[int],
actuator_biastype: wp.array(dtype=int), actuator_biastype: wp.array[int],
actuator_actadr: wp.array(dtype=int), actuator_actadr: wp.array[int],
actuator_actnum: wp.array(dtype=int), actuator_actnum: wp.array[int],
actuator_forcelimited: wp.array(dtype=bool), actuator_forcelimited: wp.array[bool],
actuator_actlimited: wp.array(dtype=bool), actuator_actlimited: wp.array[bool],
actuator_dynprm: wp.array2d(dtype=vec10f), actuator_dynprm: wp.array2d[vec10f],
actuator_gainprm: wp.array2d(dtype=vec10f), actuator_gainprm: wp.array2d[vec10f],
actuator_biasprm: wp.array2d(dtype=vec10f), actuator_biasprm: wp.array2d[vec10f],
actuator_actearly: wp.array(dtype=bool), actuator_actearly: wp.array[bool],
actuator_forcerange: wp.array2d(dtype=wp.vec2), actuator_forcerange: wp.array2d[wp.vec2],
actuator_actrange: wp.array2d(dtype=wp.vec2), actuator_actrange: wp.array2d[wp.vec2],
# Data in: # Data in:
act_in: wp.array2d(dtype=float), act_in: wp.array2d[float],
ctrl_in: wp.array2d(dtype=float), ctrl_in: wp.array2d[float],
act_dot_in: wp.array2d(dtype=float), act_dot_in: wp.array2d[float],
actuator_force_in: wp.array2d(dtype=float), actuator_force_in: wp.array2d[float],
# Out: # Out:
vel_out: wp.array2d(dtype=float), vel_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
@@ -121,16 +121,16 @@ def _qderiv_actuator_passive_actuation_dense(
# Model: # Model:
nu: int, nu: int,
# Data in: # Data in:
moment_rownnz_in: wp.array2d(dtype=int), moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d(dtype=int), moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d(dtype=int), moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d(dtype=float), actuator_moment_in: wp.array2d[float],
# In: # In:
vel_in: wp.array2d(dtype=float), vel_in: wp.array2d[float],
qMi: wp.array(dtype=int), qMi: wp.array[int],
qMj: wp.array(dtype=int), qMj: wp.array[int],
# Out: # Out:
qDeriv_out: wp.array3d(dtype=float), qDeriv_out: wp.array3d[float],
): ):
worldid, elemid = wp.tid() worldid, elemid = wp.tid()
@@ -171,18 +171,18 @@ def _qderiv_actuator_passive_actuation_dense(
@wp.kernel @wp.kernel
def _qderiv_actuator_passive_actuation_sparse( def _qderiv_actuator_passive_actuation_sparse(
# Model: # Model:
M_rownnz: wp.array(dtype=int), M_rownnz: wp.array[int],
M_rowadr: wp.array(dtype=int), M_rowadr: wp.array[int],
# Data in: # Data in:
moment_rownnz_in: wp.array2d(dtype=int), moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d(dtype=int), moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d(dtype=int), moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d(dtype=float), actuator_moment_in: wp.array2d[float],
# In: # In:
vel_in: wp.array2d(dtype=float), vel_in: wp.array2d[float],
qMj: wp.array(dtype=int), qMj: wp.array[int],
# Out: # Out:
qDeriv_out: wp.array3d(dtype=float), qDeriv_out: wp.array3d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
@@ -225,18 +225,18 @@ def _qderiv_actuator_passive_actuation_sparse(
@wp.kernel @wp.kernel
def _qderiv_actuator_passive( def _qderiv_actuator_passive(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
opt_disableflags: int, opt_disableflags: int,
dof_damping: wp.array2d(dtype=float), dof_damping: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
qM_in: wp.array3d(dtype=float), qM_in: wp.array3d[float],
# In: # In:
qMi: wp.array(dtype=int), qMi: wp.array[int],
qMj: wp.array(dtype=int), qMj: wp.array[int],
qDeriv_in: wp.array3d(dtype=float), qDeriv_in: wp.array3d[float],
# Out: # Out:
qDeriv_out: wp.array3d(dtype=float), qDeriv_out: wp.array3d[float],
): ):
worldid, elemid = wp.tid() worldid, elemid = wp.tid()
@@ -267,19 +267,19 @@ def _qderiv_actuator_passive(
def _qderiv_tendon_damping( def _qderiv_tendon_damping(
# Model: # Model:
ntendon: int, ntendon: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
tendon_damping: wp.array2d(dtype=float), tendon_damping: wp.array2d[float],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
# In: # In:
qMi: wp.array(dtype=int), qMi: wp.array[int],
qMj: wp.array(dtype=int), qMj: wp.array[int],
# Out: # Out:
qDeriv_out: wp.array3d(dtype=float), qDeriv_out: wp.array3d[float],
): ):
worldid, elemid = wp.tid() worldid, elemid = wp.tid()
dofiid = qMi[elemid] dofiid = qMi[elemid]
@@ -318,7 +318,7 @@ def _qderiv_tendon_damping(
@event_scope @event_scope
def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)): def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
"""Analytical derivative of smooth forces w.r.t. velocities. """Analytical derivative of smooth forces w.r.t. velocities.
Args: Args:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+113 -113
View File
@@ -51,17 +51,17 @@ wp.set_module_options({"enable_backward": False})
@wp.kernel @wp.kernel
def _next_position( def _next_position(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
jnt_type: wp.array(dtype=int), jnt_type: wp.array[int],
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
# Data in: # Data in:
qpos_in: wp.array2d(dtype=float), qpos_in: wp.array2d[float],
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
# In: # In:
qvel_scale_in: float, qvel_scale_in: float,
# Data out: # Data out:
qpos_out: wp.array2d(dtype=float), qpos_out: wp.array2d[float],
): ):
worldid, jntid = wp.tid() worldid, jntid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]] timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -115,14 +115,14 @@ def _next_position(
@wp.kernel @wp.kernel
def _next_velocity( def _next_velocity(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
qacc_in: wp.array2d(dtype=float), qacc_in: wp.array2d[float],
# In: # In:
qacc_scale_in: float, qacc_scale_in: float,
# Data out: # Data out:
qvel_out: wp.array2d(dtype=float), qvel_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]] timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -132,21 +132,21 @@ def _next_velocity(
@wp.kernel @wp.kernel
def _next_activation( def _next_activation(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
actuator_dyntype: wp.array(dtype=int), actuator_dyntype: wp.array[int],
actuator_actadr: wp.array(dtype=int), actuator_actadr: wp.array[int],
actuator_actnum: wp.array(dtype=int), actuator_actnum: wp.array[int],
actuator_actlimited: wp.array(dtype=bool), actuator_actlimited: wp.array[bool],
actuator_dynprm: wp.array2d(dtype=vec10f), actuator_dynprm: wp.array2d[vec10f],
actuator_actrange: wp.array2d(dtype=wp.vec2), actuator_actrange: wp.array2d[wp.vec2],
# Data in: # Data in:
act_in: wp.array2d(dtype=float), act_in: wp.array2d[float],
act_dot_in: wp.array2d(dtype=float), act_dot_in: wp.array2d[float],
# In: # In:
act_dot_scale: float, act_dot_scale: float,
limit: bool, limit: bool,
# Data out: # Data out:
act_out: wp.array2d(dtype=float), act_out: wp.array2d[float],
): ):
worldid, uid = wp.tid() worldid, uid = wp.tid()
opt_timestep_id = worldid % opt_timestep.shape[0] opt_timestep_id = worldid % opt_timestep.shape[0]
@@ -171,21 +171,21 @@ def _next_activation(
@wp.kernel @wp.kernel
def _next_time( def _next_time(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
is_sparse: bool, is_sparse: bool,
# Data in: # Data in:
nefc_in: wp.array(dtype=int), nefc_in: wp.array[int],
time_in: wp.array(dtype=float), time_in: wp.array[float],
efc_J_rownnz_in: wp.array2d(dtype=int), efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d(dtype=int), efc_J_rowadr_in: wp.array2d[int],
nworld_in: int, nworld_in: int,
naconmax_in: int, naconmax_in: int,
njmax_in: int, njmax_in: int,
njmax_nnz_in: int, njmax_nnz_in: int,
nacon_in: wp.array(dtype=int), nacon_in: wp.array[int],
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array[int],
# Data out: # Data out:
time_out: wp.array(dtype=float), time_out: wp.array[float],
): ):
worldid = wp.tid() worldid = wp.tid()
time_out[worldid] = time_in[worldid] + opt_timestep[worldid % opt_timestep.shape[0]] time_out[worldid] = time_in[worldid] + opt_timestep[worldid % opt_timestep.shape[0]]
@@ -277,11 +277,11 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
@wp.kernel @wp.kernel
def _euler_damp_qfrc_sparse( def _euler_damp_qfrc_sparse(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
dof_Madr: wp.array(dtype=int), dof_Madr: wp.array[int],
dof_damping: wp.array2d(dtype=float), dof_damping: wp.array2d[float],
# Out: # Out:
qM_integration_out: wp.array3d(dtype=float), qM_integration_out: wp.array3d[float],
): ):
worldid, tid = wp.tid() worldid, tid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]] timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -295,15 +295,15 @@ def _tile_euler_dense(tile: TileSet):
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def euler_dense( def euler_dense(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
dof_damping: wp.array2d(dtype=float), dof_damping: wp.array2d[float],
# Data in: # Data in:
qM_in: wp.array3d(dtype=float), qM_in: wp.array3d[float],
efc_Ma_in: wp.array2d(dtype=float), efc_Ma_in: wp.array2d[float],
# In: # In:
adr_in: wp.array(dtype=int), adr_in: wp.array[int],
# Data out: # Data out:
qacc_out: wp.array2d(dtype=float), qacc_out: wp.array2d[float],
): ):
worldid, nodeid = wp.tid() worldid, nodeid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]] timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -358,8 +358,8 @@ def _rk_perturb_state(
m: Model, m: Model,
d: Data, d: Data,
scale: float, scale: float,
qpos_t0: wp.array2d(dtype=float), qpos_t0: wp.array2d[float],
qvel_t0: wp.array2d(dtype=float), qvel_t0: wp.array2d[float],
act_t0: Optional[wp.array] = None, act_t0: Optional[wp.array] = None,
): ):
# position # position
@@ -403,13 +403,13 @@ def _rk_perturb_state(
@wp.kernel @wp.kernel
def _rk_accumulate_velocity_acceleration( def _rk_accumulate_velocity_acceleration(
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
qacc_in: wp.array2d(dtype=float), qacc_in: wp.array2d[float],
# In: # In:
scale: float, scale: float,
# Data out: # Data out:
qvel_out: wp.array2d(dtype=float), qvel_out: wp.array2d[float],
qacc_out: wp.array2d(dtype=float), qacc_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
qvel_out[worldid, dofid] += scale * qvel_in[worldid, dofid] qvel_out[worldid, dofid] += scale * qvel_in[worldid, dofid]
@@ -419,11 +419,11 @@ def _rk_accumulate_velocity_acceleration(
@wp.kernel @wp.kernel
def _rk_accumulate_activation_velocity( def _rk_accumulate_activation_velocity(
# Data in: # Data in:
act_dot_in: wp.array2d(dtype=float), act_dot_in: wp.array2d[float],
# In: # In:
scale: float, scale: float,
# Data out: # Data out:
act_dot_out: wp.array2d(dtype=float), act_dot_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
act_dot_out[worldid, actid] += scale * act_dot_in[worldid, actid] act_dot_out[worldid, actid] += scale * act_dot_in[worldid, actid]
@@ -433,8 +433,8 @@ def _rk_accumulate(
m: Model, m: Model,
d: Data, d: Data,
scale: float, scale: float,
qvel_rk: wp.array2d(dtype=float), qvel_rk: wp.array2d[float],
qacc_rk: wp.array2d(dtype=float), qacc_rk: wp.array2d[float],
act_dot_rk: Optional[wp.array] = None, act_dot_rk: Optional[wp.array] = None,
): ):
"""Computes one term of 1/6 k_1 + 1/3 k_2 + 1/3 k_3 + 1/6 k_4.""" """Computes one term of 1/6 k_1 + 1/3 k_2 + 1/3 k_3 + 1/6 k_4."""
@@ -540,13 +540,13 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
@wp.kernel @wp.kernel
def _actuator_velocity( def _actuator_velocity(
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
moment_rownnz_in: wp.array2d(dtype=int), moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d(dtype=int), moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d(dtype=int), moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d(dtype=float), actuator_moment_in: wp.array2d[float],
# Data out: # Data out:
actuator_velocity_out: wp.array2d(dtype=float), actuator_velocity_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
@@ -565,14 +565,14 @@ def _actuator_velocity(
@wp.kernel @wp.kernel
def _tendon_velocity( def _tendon_velocity(
# Model: # Model:
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
# Data in: # Data in:
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
# Data out: # Data out:
ten_velocity_out: wp.array2d(dtype=float), ten_velocity_out: wp.array2d[float],
): ):
worldid, tenid = wp.tid() worldid, tenid = wp.tid()
@@ -617,34 +617,34 @@ def fwd_velocity(m: Model, d: Data):
def _actuator_force( def _actuator_force(
# Model: # Model:
na: int, na: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
actuator_dyntype: wp.array(dtype=int), actuator_dyntype: wp.array[int],
actuator_gaintype: wp.array(dtype=int), actuator_gaintype: wp.array[int],
actuator_biastype: wp.array(dtype=int), actuator_biastype: wp.array[int],
actuator_actadr: wp.array(dtype=int), actuator_actadr: wp.array[int],
actuator_actnum: wp.array(dtype=int), actuator_actnum: wp.array[int],
actuator_ctrllimited: wp.array(dtype=bool), actuator_ctrllimited: wp.array[bool],
actuator_forcelimited: wp.array(dtype=bool), actuator_forcelimited: wp.array[bool],
actuator_actlimited: wp.array(dtype=bool), actuator_actlimited: wp.array[bool],
actuator_dynprm: wp.array2d(dtype=vec10f), actuator_dynprm: wp.array2d[vec10f],
actuator_gainprm: wp.array2d(dtype=vec10f), actuator_gainprm: wp.array2d[vec10f],
actuator_biasprm: wp.array2d(dtype=vec10f), actuator_biasprm: wp.array2d[vec10f],
actuator_actearly: wp.array(dtype=bool), actuator_actearly: wp.array[bool],
actuator_ctrlrange: wp.array2d(dtype=wp.vec2), actuator_ctrlrange: wp.array2d[wp.vec2],
actuator_forcerange: wp.array2d(dtype=wp.vec2), actuator_forcerange: wp.array2d[wp.vec2],
actuator_actrange: wp.array2d(dtype=wp.vec2), actuator_actrange: wp.array2d[wp.vec2],
actuator_acc0: wp.array2d(dtype=float), actuator_acc0: wp.array2d[float],
actuator_lengthrange: wp.array2d(dtype=wp.vec2), actuator_lengthrange: wp.array2d[wp.vec2],
# Data in: # Data in:
act_in: wp.array2d(dtype=float), act_in: wp.array2d[float],
ctrl_in: wp.array2d(dtype=float), ctrl_in: wp.array2d[float],
actuator_length_in: wp.array2d(dtype=float), actuator_length_in: wp.array2d[float],
actuator_velocity_in: wp.array2d(dtype=float), actuator_velocity_in: wp.array2d[float],
# In: # In:
dsbl_clampctrl: int, dsbl_clampctrl: int,
# Data out: # Data out:
act_dot_out: wp.array2d(dtype=float), act_dot_out: wp.array2d[float],
actuator_force_out: wp.array2d(dtype=float), actuator_force_out: wp.array2d[float],
): ):
worldid, uid = wp.tid() worldid, uid = wp.tid()
@@ -738,12 +738,12 @@ def _actuator_force(
@wp.kernel @wp.kernel
def _tendon_actuator_force( def _tendon_actuator_force(
# Model: # Model:
actuator_trntype: wp.array(dtype=int), actuator_trntype: wp.array[int],
actuator_trnid: wp.array(dtype=wp.vec2i), actuator_trnid: wp.array[wp.vec2i],
# Data in: # Data in:
actuator_force_in: wp.array2d(dtype=float), actuator_force_in: wp.array2d[float],
# Out: # Out:
ten_actfrc_out: wp.array2d(dtype=float), ten_actfrc_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
@@ -756,14 +756,14 @@ def _tendon_actuator_force(
@wp.kernel @wp.kernel
def _tendon_actuator_force_clamp( def _tendon_actuator_force_clamp(
# Model: # Model:
tendon_actfrclimited: wp.array(dtype=bool), tendon_actfrclimited: wp.array[bool],
tendon_actfrcrange: wp.array2d(dtype=wp.vec2), tendon_actfrcrange: wp.array2d[wp.vec2],
actuator_trntype: wp.array(dtype=int), actuator_trntype: wp.array[int],
actuator_trnid: wp.array(dtype=wp.vec2i), actuator_trnid: wp.array[wp.vec2i],
# In: # In:
ten_actfrc_in: wp.array2d(dtype=float), ten_actfrc_in: wp.array2d[float],
# Data out: # Data out:
actuator_force_out: wp.array2d(dtype=float), actuator_force_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
@@ -782,13 +782,13 @@ def _tendon_actuator_force_clamp(
@wp.kernel @wp.kernel
def _qfrc_actuator( def _qfrc_actuator(
# Data in: # Data in:
moment_rownnz_in: wp.array2d(dtype=int), moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d(dtype=int), moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d(dtype=int), moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d(dtype=float), actuator_moment_in: wp.array2d[float],
actuator_force_in: wp.array2d(dtype=float), actuator_force_in: wp.array2d[float],
# Data out: # Data out:
qfrc_actuator_out: wp.array2d(dtype=float), qfrc_actuator_out: wp.array2d[float],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
@@ -806,15 +806,15 @@ def _qfrc_actuator(
def _qfrc_actuator_gravcomp_limits( def _qfrc_actuator_gravcomp_limits(
# Model: # Model:
ngravcomp: int, ngravcomp: int,
jnt_actfrclimited: wp.array(dtype=bool), jnt_actfrclimited: wp.array[bool],
jnt_actgravcomp: wp.array(dtype=int), jnt_actgravcomp: wp.array[int],
jnt_actfrcrange: wp.array2d(dtype=wp.vec2), jnt_actfrcrange: wp.array2d[wp.vec2],
dof_jntid: wp.array(dtype=int), dof_jntid: wp.array[int],
# Data in: # Data in:
qfrc_gravcomp_in: wp.array2d(dtype=float), qfrc_gravcomp_in: wp.array2d[float],
qfrc_actuator_in: wp.array2d(dtype=float), qfrc_actuator_in: wp.array2d[float],
# Data out: # Data out:
qfrc_actuator_out: wp.array2d(dtype=float), qfrc_actuator_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
jntid = dof_jntid[dofid] jntid = dof_jntid[dofid]
@@ -930,12 +930,12 @@ def fwd_actuation(m: Model, d: Data):
@wp.kernel @wp.kernel
def _qfrc_smooth( def _qfrc_smooth(
# Data in: # Data in:
qfrc_applied_in: wp.array2d(dtype=float), qfrc_applied_in: wp.array2d[float],
qfrc_bias_in: wp.array2d(dtype=float), qfrc_bias_in: wp.array2d[float],
qfrc_passive_in: wp.array2d(dtype=float), qfrc_passive_in: wp.array2d[float],
qfrc_actuator_in: wp.array2d(dtype=float), qfrc_actuator_in: wp.array2d[float],
# Data out: # Data out:
qfrc_smooth_out: wp.array2d(dtype=float), qfrc_smooth_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
qfrc_smooth_out[worldid, dofid] = ( qfrc_smooth_out[worldid, dofid] = (
mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
Vendored
+10 -10
View File
@@ -34,12 +34,12 @@ wp.set_module_options({"enable_backward": False})
@wp.kernel @wp.kernel
def _qfrc_eulerdamp( def _qfrc_eulerdamp(
# Model: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
dof_damping: wp.array2d(dtype=float), dof_damping: wp.array2d[float],
# Data in: # Data in:
qacc_in: wp.array2d(dtype=float), qacc_in: wp.array2d[float],
# Out: # Out:
qfrc_out: wp.array2d(dtype=float), qfrc_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]] timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -49,13 +49,13 @@ def _qfrc_eulerdamp(
@wp.kernel @wp.kernel
def _qfrc_inverse( def _qfrc_inverse(
# Data in: # Data in:
qfrc_bias_in: wp.array2d(dtype=float), qfrc_bias_in: wp.array2d[float],
qfrc_passive_in: wp.array2d(dtype=float), qfrc_passive_in: wp.array2d[float],
qfrc_constraint_in: wp.array2d(dtype=float), qfrc_constraint_in: wp.array2d[float],
# In: # In:
Ma: wp.array2d(dtype=float), Ma: wp.array2d[float],
# Data out: # Data out:
qfrc_inverse_out: wp.array2d(dtype=float), qfrc_inverse_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
@@ -67,7 +67,7 @@ def _qfrc_inverse(
qfrc_inverse_out[worldid, dofid] = qfrc_inverse qfrc_inverse_out[worldid, dofid] = qfrc_inverse
def discrete_acc(m: Model, d: Data, qacc: wp.array2d(dtype=float)): def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
"""Convert discrete-time qacc to continuous-time qacc. """Convert discrete-time qacc to continuous-time qacc.
Args: Args:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+161 -151
View File
@@ -34,14 +34,20 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array | None: def _is_array_spec(typ) -> bool:
"""Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
return isinstance(typ, wp.array) or type(typ).__name__ == "_ArrayAnnotation"
def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
"""Creates a warp array and populates it with data. """Creates a warp array and populates it with data.
The array shape is determined by a field spec referencing MjModel / MjData array sizes. The array shape is determined by a field spec referencing MjModel / MjData array sizes.
""" """
spec_shape = getattr(spec, "shape", (0,))
shape = None shape = None
if spec.shape != (0,): if spec_shape != (0,):
shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec.shape) shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec_shape)
if data is None and shape is None: if data is None and shape is None:
return None # nothing to do return None # nothing to do
@@ -50,7 +56,7 @@ def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array
else: else:
array = wp.array(np.array(data), dtype=spec.dtype, shape=shape) array = wp.array(np.array(data), dtype=spec.dtype, shape=shape)
if spec.shape[0] == "*": if spec_shape and spec_shape[0] == "*":
# add private attribute for JAX to determine which fields are batched # add private attribute for JAX to determine which fields are batched
array._is_batched = True array._is_batched = True
# also set stride 0 to 0 which is expected legacy behavior (but is deprecated) # also set stride 0 to 0 which is expected legacy behavior (but is deprecated)
@@ -194,7 +200,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
# place opt on device # place opt on device
for f in dataclasses.fields(types.Option): for f in dataclasses.fields(types.Option):
if isinstance(f.type, wp.array): if _is_array_spec(f.type):
setattr(opt, f.name, _create_array(getattr(opt, f.name), f.type, {"*": 1})) setattr(opt, f.name, _create_array(getattr(opt, f.name), f.type, {"*": 1}))
else: else:
setattr(opt, f.name, f.type(getattr(opt, f.name))) setattr(opt, f.name, f.type(getattr(opt, f.name)))
@@ -635,7 +641,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
# place m on device # place m on device
sizes = dict({"*": 1}, **{f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int}) sizes = dict({"*": 1}, **{f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int})
for f in dataclasses.fields(types.Model): for f in dataclasses.fields(types.Model):
if isinstance(f.type, wp.array): if _is_array_spec(f.type):
setattr(m, f.name, _create_array(getattr(m, f.name), f.type, sizes)) setattr(m, f.name, _create_array(getattr(m, f.name), f.type, sizes))
return m return m
@@ -1459,7 +1465,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
""" """
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def reset_xfrc_applied(reset_in: wp.array(dtype=bool), xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector)): def reset_xfrc_applied(reset_in: wp.array[bool], xfrc_applied_out: wp.array2d[wp.spatial_vector]):
worldid, bodyid, elemid = wp.tid() worldid, bodyid, elemid = wp.tid()
if wp.static(reset is not None): if wp.static(reset is not None):
@@ -1469,7 +1475,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
xfrc_applied_out[worldid, bodyid][elemid] = 0.0 xfrc_applied_out[worldid, bodyid][elemid] = 0.0
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def reset_qM(reset_in: wp.array(dtype=bool), qM_out: wp.array3d(dtype=float)): def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
worldid, elemid1, elemid2 = wp.tid() worldid, elemid1, elemid2 = wp.tid()
if wp.static(reset is not None): if wp.static(reset is not None):
@@ -1487,31 +1493,31 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
na: int, na: int,
neq: int, neq: int,
nsensordata: int, nsensordata: int,
qpos0: wp.array2d(dtype=float), qpos0: wp.array2d[float],
eq_active0: wp.array(dtype=bool), eq_active0: wp.array[bool],
# Data in: # Data in:
nworld_in: int, nworld_in: int,
# In: # In:
reset_in: wp.array(dtype=bool), reset_in: wp.array[bool],
# Data out: # Data out:
solver_niter_out: wp.array(dtype=int), solver_niter_out: wp.array[int],
ne_out: wp.array(dtype=int), ne_out: wp.array[int],
nf_out: wp.array(dtype=int), nf_out: wp.array[int],
nl_out: wp.array(dtype=int), nl_out: wp.array[int],
nefc_out: wp.array(dtype=int), nefc_out: wp.array[int],
time_out: wp.array(dtype=float), time_out: wp.array[float],
energy_out: wp.array(dtype=wp.vec2), energy_out: wp.array[wp.vec2],
qpos_out: wp.array2d(dtype=float), qpos_out: wp.array2d[float],
qvel_out: wp.array2d(dtype=float), qvel_out: wp.array2d[float],
act_out: wp.array2d(dtype=float), act_out: wp.array2d[float],
qacc_warmstart_out: wp.array2d(dtype=float), qacc_warmstart_out: wp.array2d[float],
ctrl_out: wp.array2d(dtype=float), ctrl_out: wp.array2d[float],
qfrc_applied_out: wp.array2d(dtype=float), qfrc_applied_out: wp.array2d[float],
eq_active_out: wp.array2d(dtype=bool), eq_active_out: wp.array2d[bool],
qacc_out: wp.array2d(dtype=float), qacc_out: wp.array2d[float],
act_dot_out: wp.array2d(dtype=float), act_dot_out: wp.array2d[float],
sensordata_out: wp.array2d(dtype=float), sensordata_out: wp.array2d[float],
nacon_out: wp.array(dtype=int), nacon_out: wp.array[int],
): ):
worldid = wp.tid() worldid = wp.tid()
@@ -1549,14 +1555,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def reset_mocap( def reset_mocap(
# Model: # Model:
body_mocapid: wp.array(dtype=int), body_mocapid: wp.array[int],
body_pos: wp.array2d(dtype=wp.vec3), body_pos: wp.array2d[wp.vec3],
body_quat: wp.array2d(dtype=wp.quat), body_quat: wp.array2d[wp.quat],
# In: # In:
reset_in: wp.array(dtype=bool), reset_in: wp.array[bool],
# Data out: # Data out:
mocap_pos_out: wp.array2d(dtype=wp.vec3), mocap_pos_out: wp.array2d[wp.vec3],
mocap_quat_out: wp.array2d(dtype=wp.quat), mocap_quat_out: wp.array2d[wp.quat],
): ):
worldid, bodyid = wp.tid() worldid, bodyid = wp.tid()
@@ -1573,27 +1579,27 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def reset_contact( def reset_contact(
# Data in: # Data in:
nacon_in: wp.array(dtype=int), nacon_in: wp.array[int],
# In: # In:
reset_in: wp.array(dtype=bool), reset_in: wp.array[bool],
nefcaddress: int, nefcaddress: int,
# Data out: # Data out:
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array[float],
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array[wp.vec3],
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array[wp.mat33],
contact_includemargin_out: wp.array(dtype=float), contact_includemargin_out: wp.array[float],
contact_friction_out: wp.array(dtype=types.vec5), contact_friction_out: wp.array[types.vec5],
contact_solref_out: wp.array(dtype=wp.vec2), contact_solref_out: wp.array[wp.vec2],
contact_solreffriction_out: wp.array(dtype=wp.vec2), contact_solreffriction_out: wp.array[wp.vec2],
contact_solimp_out: wp.array(dtype=types.vec5), contact_solimp_out: wp.array[types.vec5],
contact_dim_out: wp.array(dtype=int), contact_dim_out: wp.array[int],
contact_geom_out: wp.array(dtype=wp.vec2i), contact_geom_out: wp.array[wp.vec2i],
contact_flex_out: wp.array(dtype=wp.vec2i), contact_flex_out: wp.array[wp.vec2i],
contact_vert_out: wp.array(dtype=wp.vec2i), contact_vert_out: wp.array[wp.vec2i],
contact_efc_address_out: wp.array2d(dtype=int), contact_efc_address_out: wp.array2d[int],
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array[int],
contact_type_out: wp.array(dtype=int), contact_type_out: wp.array[int],
contact_geomcollisionid_out: wp.array(dtype=int), contact_geomcollisionid_out: wp.array[int],
): ):
conid = wp.tid() conid = wp.tid()
@@ -1697,8 +1703,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
# kernel_analyzer: off # kernel_analyzer: off
@wp.kernel @wp.kernel
def _init_subtreemass( def _init_subtreemass(
body_mass_in: wp.array2d(dtype=float), body_mass_in: wp.array2d[float],
body_subtreemass_out: wp.array2d(dtype=float), body_subtreemass_out: wp.array2d[float],
): ):
worldid, bodyid = wp.tid() worldid, bodyid = wp.tid()
body_mass_id = worldid % body_mass_in.shape[0] body_mass_id = worldid % body_mass_in.shape[0]
@@ -1708,9 +1714,9 @@ def _init_subtreemass(
@wp.kernel @wp.kernel
def _accumulate_subtreemass( def _accumulate_subtreemass(
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_subtreemass_io: wp.array2d(dtype=float), body_subtreemass_io: wp.array2d[float],
body_tree_: wp.array(dtype=int), body_tree_: wp.array[int],
): ):
worldid, nodeid = wp.tid() worldid, nodeid = wp.tid()
body_subtreemass_id = worldid % body_subtreemass_io.shape[0] body_subtreemass_id = worldid % body_subtreemass_io.shape[0]
@@ -1722,8 +1728,8 @@ def _accumulate_subtreemass(
@wp.kernel @wp.kernel
def _copy_qpos0_to_qpos( def _copy_qpos0_to_qpos(
qpos0: wp.array2d(dtype=float), qpos0: wp.array2d[float],
qpos_out: wp.array2d(dtype=float), qpos_out: wp.array2d[float],
): ):
worldid, i = wp.tid() worldid, i = wp.tid()
qpos0_id = worldid % qpos0.shape[0] qpos0_id = worldid % qpos0.shape[0]
@@ -1732,8 +1738,8 @@ def _copy_qpos0_to_qpos(
@wp.kernel @wp.kernel
def _copy_tendon_length0( def _copy_tendon_length0(
ten_length_in: wp.array2d(dtype=float), ten_length_in: wp.array2d[float],
tendon_length0_out: wp.array2d(dtype=float), tendon_length0_out: wp.array2d[float],
): ):
worldid, tenid = wp.tid() worldid, tenid = wp.tid()
tendon_length0_id = worldid % tendon_length0_out.shape[0] tendon_length0_id = worldid % tendon_length0_out.shape[0]
@@ -1744,9 +1750,9 @@ def _copy_tendon_length0(
def _compute_meaninertia( def _compute_meaninertia(
nv: int, nv: int,
is_sparse: bool, is_sparse: bool,
dof_Madr_in: wp.array(dtype=int), dof_Madr_in: wp.array[int],
qM_in: wp.array3d(dtype=float), qM_in: wp.array3d[float],
meaninertia_out: wp.array(dtype=float), meaninertia_out: wp.array[float],
): ):
"""Compute mean diagonal inertia from qM at qpos0.""" """Compute mean diagonal inertia from qM at qpos0."""
worldid = wp.tid() worldid = wp.tid()
@@ -1771,7 +1777,7 @@ def _compute_meaninertia(
@wp.kernel @wp.kernel
def _set_unit_vector( def _set_unit_vector(
dofid_target: int, dofid_target: int,
unit_vec_out: wp.array2d(dtype=float), unit_vec_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
nv = unit_vec_out.shape[1] nv = unit_vec_out.shape[1]
@@ -1785,8 +1791,8 @@ def _set_unit_vector(
@wp.kernel @wp.kernel
def _extract_dof_A_diag( def _extract_dof_A_diag(
dofid: int, dofid: int,
result_vec_in: wp.array2d(dtype=float), result_vec_in: wp.array2d[float],
dof_A_diag_out: wp.array2d(dtype=float), dof_A_diag_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
dof_A_diag_id = worldid % dof_A_diag_out.shape[0] dof_A_diag_id = worldid % dof_A_diag_out.shape[0]
@@ -1795,11 +1801,11 @@ def _extract_dof_A_diag(
@wp.kernel @wp.kernel
def _finalize_dof_invweight0( def _finalize_dof_invweight0(
dof_jntid: wp.array(dtype=int), dof_jntid: wp.array[int],
jnt_type: wp.array(dtype=int), jnt_type: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
dof_A_diag_in: wp.array2d(dtype=float), dof_A_diag_in: wp.array2d[float],
dof_invweight0_out: wp.array2d(dtype=float), dof_invweight0_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
dof_invweight0_id = worldid % dof_invweight0_out.shape[0] dof_invweight0_id = worldid % dof_invweight0_out.shape[0]
@@ -1842,15 +1848,15 @@ def _compute_body_jac_row(
nv: int, nv: int,
bodyid_target: int, bodyid_target: int,
row_idx: int, row_idx: int,
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array[int],
dof_parentid: wp.array(dtype=int), dof_parentid: wp.array[int],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
xipos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
body_jac_row_out: wp.array2d(dtype=float), body_jac_row_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
@@ -1901,9 +1907,9 @@ def _compute_body_A_diag_entry(
nv: int, nv: int,
bodyid_target: int, bodyid_target: int,
row_idx: int, row_idx: int,
body_jac_row_in: wp.array2d(dtype=float), body_jac_row_in: wp.array2d[float],
result_vec_in: wp.array2d(dtype=float), result_vec_in: wp.array2d[float],
body_A_diag_out: wp.array3d(dtype=float), body_A_diag_out: wp.array3d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
body_A_diag_id = worldid % body_A_diag_out.shape[0] body_A_diag_id = worldid % body_A_diag_out.shape[0]
@@ -1916,9 +1922,9 @@ def _compute_body_A_diag_entry(
@wp.kernel @wp.kernel
def _finalize_body_invweight0( def _finalize_body_invweight0(
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
body_A_diag_in: wp.array3d(dtype=float), body_A_diag_in: wp.array3d[float],
body_invweight0_out: wp.array2d(dtype=wp.vec2), body_invweight0_out: wp.array2d[wp.vec2],
): ):
worldid, bodyid = wp.tid() worldid, bodyid = wp.tid()
body_invweight0_id = worldid % body_invweight0_out.shape[0] body_invweight0_id = worldid % body_invweight0_out.shape[0]
@@ -1953,11 +1959,11 @@ def _finalize_body_invweight0(
@wp.kernel @wp.kernel
def _copy_tendon_jacobian( def _copy_tendon_jacobian(
tenid_target: int, tenid_target: int,
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
ten_J_vec_out: wp.array2d(dtype=float), ten_J_vec_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
nv = ten_J_in.shape[2] nv = ten_J_in.shape[2]
@@ -1971,15 +1977,15 @@ def _copy_tendon_jacobian(
@wp.kernel @wp.kernel
def _compute_tendon_dot_product( def _compute_tendon_dot_product(
# Model: # Model:
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
# In: # In:
tenid_target: int, tenid_target: int,
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
result_vec_in: wp.array2d(dtype=float), result_vec_in: wp.array2d[float],
# Out: # Out:
tendon_invweight0_out: wp.array2d(dtype=float), tendon_invweight0_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
tendon_invweight0_id = worldid % tendon_invweight0_out.shape[0] tendon_invweight0_id = worldid % tendon_invweight0_out.shape[0]
@@ -1997,15 +2003,15 @@ def _compute_tendon_dot_product(
@wp.kernel @wp.kernel
def _compute_cam_pos0( def _compute_cam_pos0(
cam_bodyid: wp.array(dtype=int), cam_bodyid: wp.array[int],
cam_targetbodyid: wp.array(dtype=int), cam_targetbodyid: wp.array[int],
cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xpos_in: wp.array2d[wp.vec3],
cam_xmat_in: wp.array2d(dtype=wp.mat33), cam_xmat_in: wp.array2d[wp.mat33],
xpos_in: wp.array2d(dtype=wp.vec3), xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cam_pos0_out: wp.array2d(dtype=wp.vec3), cam_pos0_out: wp.array2d[wp.vec3],
cam_poscom0_out: wp.array2d(dtype=wp.vec3), cam_poscom0_out: wp.array2d[wp.vec3],
cam_mat0_out: wp.array2d(dtype=wp.mat33), cam_mat0_out: wp.array2d[wp.mat33],
): ):
worldid, camid = wp.tid() worldid, camid = wp.tid()
cam_pos0_id = worldid % cam_pos0_out.shape[0] cam_pos0_id = worldid % cam_pos0_out.shape[0]
@@ -2023,15 +2029,15 @@ def _compute_cam_pos0(
@wp.kernel @wp.kernel
def _compute_light_pos0( def _compute_light_pos0(
light_bodyid: wp.array(dtype=int), light_bodyid: wp.array[int],
light_targetbodyid: wp.array(dtype=int), light_targetbodyid: wp.array[int],
light_xpos_in: wp.array2d(dtype=wp.vec3), light_xpos_in: wp.array2d[wp.vec3],
light_xdir_in: wp.array2d(dtype=wp.vec3), light_xdir_in: wp.array2d[wp.vec3],
xpos_in: wp.array2d(dtype=wp.vec3), xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
light_pos0_out: wp.array2d(dtype=wp.vec3), light_pos0_out: wp.array2d[wp.vec3],
light_poscom0_out: wp.array2d(dtype=wp.vec3), light_poscom0_out: wp.array2d[wp.vec3],
light_dir0_out: wp.array2d(dtype=wp.vec3), light_dir0_out: wp.array2d[wp.vec3],
): ):
worldid, lightid = wp.tid() worldid, lightid = wp.tid()
light_pos0_id = worldid % light_pos0_out.shape[0] light_pos0_id = worldid % light_pos0_out.shape[0]
@@ -2050,11 +2056,11 @@ def _compute_light_pos0(
@wp.kernel @wp.kernel
def _copy_actuator_moment( def _copy_actuator_moment(
actid_target: int, actid_target: int,
moment_rownnz_in: wp.array2d(dtype=int), moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d(dtype=int), moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d(dtype=int), moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d(dtype=float), actuator_moment_in: wp.array2d[float],
act_moment_vec_out: wp.array2d(dtype=float), act_moment_vec_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
nv = act_moment_vec_out.shape[1] nv = act_moment_vec_out.shape[1]
@@ -2072,8 +2078,8 @@ def _copy_actuator_moment(
def _compute_actuator_acc0( def _compute_actuator_acc0(
actid_target: int, actid_target: int,
nv: int, nv: int,
result_vec_in: wp.array2d(dtype=float), result_vec_in: wp.array2d[float],
actuator_acc0_out: wp.array2d(dtype=float), actuator_acc0_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
norm_sq = float(0.0) norm_sq = float(0.0)
@@ -2084,11 +2090,11 @@ def _compute_actuator_acc0(
@wp.kernel @wp.kernel
def _compute_dof_M0( def _compute_dof_M0(
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
dof_armature: wp.array2d(dtype=float), dof_armature: wp.array2d[float],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
crb_in: wp.array2d(dtype=vec10), crb_in: wp.array2d[vec10],
dof_M0_out: wp.array2d(dtype=float), dof_M0_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
bodyid = dof_bodyid[dofid] bodyid = dof_bodyid[dofid]
@@ -2099,15 +2105,15 @@ def _compute_dof_M0(
@wp.kernel @wp.kernel
def _resolve_dampratio( def _resolve_dampratio(
actuator_biastype: wp.array(dtype=int), actuator_biastype: wp.array[int],
actuator_gainprm: wp.array2d(dtype=types.vec10f), actuator_gainprm: wp.array2d[types.vec10f],
moment_rownnz_in: wp.array2d(dtype=int), moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d(dtype=int), moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d(dtype=int), moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d(dtype=float), actuator_moment_in: wp.array2d[float],
dof_M0_in: wp.array2d(dtype=float), dof_M0_in: wp.array2d[float],
nv: int, nv: int,
actuator_biasprm: wp.array2d(dtype=types.vec10f), actuator_biasprm: wp.array2d[types.vec10f],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
biastype = actuator_biastype[actid] biastype = actuator_biastype[actid]
@@ -2150,15 +2156,15 @@ def _resolve_dampratio(
@wp.kernel @wp.kernel
def _set_length_range( def _set_length_range(
actuator_trntype: wp.array(dtype=int), actuator_trntype: wp.array[int],
actuator_trnid: wp.array(dtype=wp.vec2i), actuator_trnid: wp.array[wp.vec2i],
actuator_gear: wp.array2d(dtype=wp.spatial_vector), actuator_gear: wp.array2d[wp.spatial_vector],
jnt_limited: wp.array(dtype=int), jnt_limited: wp.array[int],
jnt_range: wp.array2d(dtype=wp.vec2), jnt_range: wp.array2d[wp.vec2],
tendon_limited: wp.array(dtype=int), tendon_limited: wp.array[int],
tendon_range: wp.array2d(dtype=wp.vec2), tendon_range: wp.array2d[wp.vec2],
ntendon: int, ntendon: int,
actuator_lengthrange_out: wp.array2d(dtype=wp.vec2), actuator_lengthrange_out: wp.array2d[wp.vec2],
): ):
worldid, actid = wp.tid() worldid, actid = wp.tid()
trntype = actuator_trntype[actid] trntype = actuator_trntype[actid]
@@ -2515,7 +2521,13 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
"AUTO": mujoco.mjtJacobian.mjJAC_AUTO, "AUTO": mujoco.mjtJacobian.mjJAC_AUTO,
}, },
} }
mjw_only_fields = {"opt.broadphase", "opt.broadphase_filter", "opt.ls_parallel", "opt.graph_conditional"} mjw_only_fields = {
"opt.broadphase",
"opt.broadphase_filter",
"opt.ls_parallel",
"opt.graph_conditional",
"opt.contact_sensor_maxmatch",
}
mj_only_fields = {"opt.jacobian"} mj_only_fields = {"opt.jacobian"}
if not isinstance(overrides, dict): if not isinstance(overrides, dict):
@@ -2626,7 +2638,7 @@ def _build_rays(
intrinsic: wp.vec4, intrinsic: wp.vec4,
znear: float, znear: float,
# Out: # Out:
ray_out: wp.array(dtype=wp.vec3), ray_out: wp.array[wp.vec3],
): ):
xid, yid = wp.tid() xid, yid = wp.tid()
ray_out[offset + xid + yid * img_w] = render_util.compute_ray( ray_out[offset + xid + yid * img_w] = render_util.compute_ray(
@@ -2682,18 +2694,16 @@ def create_render_context(
if callable(_cubql_avail) and _cubql_avail(): if callable(_cubql_avail) and _cubql_avail():
constructor = "cubql" constructor = "cubql"
# Mesh BVHs # Mesh BVHs – build for all meshes so per-world variants are available
nmesh = mjm.nmesh nmesh = mjm.nmesh
geom_enabled_mask = np.isin(mjm.geom_group, list(enabled_geom_groups)) geom_enabled_mask = np.isin(mjm.geom_group, list(enabled_geom_groups))
mesh_geom_mask = geom_enabled_mask & (mjm.geom_type == types.GeomType.MESH) & (mjm.geom_dataid >= 0)
used_mesh_id = set(mjm.geom_dataid[mesh_geom_mask].astype(int))
geom_enabled_idx = np.nonzero(geom_enabled_mask)[0] geom_enabled_idx = np.nonzero(geom_enabled_mask)[0]
mesh_registry = {} mesh_registry = {}
mesh_bvh_id = [wp.uint64(0) for _ in range(nmesh)] mesh_bvh_id = [wp.uint64(0) for _ in range(nmesh)]
mesh_bounds_size = [wp.vec3(0.0, 0.0, 0.0) for _ in range(nmesh)] mesh_bounds_size = [wp.vec3(0.0, 0.0, 0.0) for _ in range(nmesh)]
for mid in used_mesh_id: for mid in range(nmesh):
mesh, half = bvh.build_mesh_bvh(mjm, mid, constructor=constructor) mesh, half = bvh.build_mesh_bvh(mjm, mid, constructor=constructor)
mesh_registry[mesh.id] = mesh mesh_registry[mesh.id] = mesh
mesh_bvh_id[mid] = mesh.id mesh_bvh_id[mid] = mesh.id
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+22 -22
View File
@@ -26,24 +26,24 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
def _tree_edges( def _tree_edges(
# Model: # Model:
nv: int, nv: int,
body_treeid: wp.array(dtype=int), body_treeid: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
dof_treeid: wp.array(dtype=int), dof_treeid: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
site_bodyid: wp.array(dtype=int), site_bodyid: wp.array[int],
eq_type: wp.array(dtype=int), eq_type: wp.array[int],
eq_obj1id: wp.array(dtype=int), eq_obj1id: wp.array[int],
eq_obj2id: wp.array(dtype=int), eq_obj2id: wp.array[int],
eq_objtype: wp.array(dtype=int), eq_objtype: wp.array[int],
# Data in: # Data in:
nefc_in: wp.array(dtype=int), nefc_in: wp.array[int],
contact_geom_in: wp.array(dtype=wp.vec2i), contact_geom_in: wp.array[wp.vec2i],
efc_type_in: wp.array2d(dtype=int), efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d(dtype=int), efc_id_in: wp.array2d[int],
efc_J_in: wp.array3d(dtype=float), efc_J_in: wp.array3d[float],
njmax_in: int, njmax_in: int,
# Out: # Out:
tree_tree: wp.array3d(dtype=int), # kernel_analyzer: off tree_tree: wp.array3d[int], # kernel_analyzer: off
): ):
"""Find tree edges.""" """Find tree edges."""
worldid, efcid = wp.tid() worldid, efcid = wp.tid()
@@ -151,7 +151,7 @@ def _tree_edges(
wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1) wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)): def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
"""Compute tree-tree adjacency matrix.""" """Compute tree-tree adjacency matrix."""
tree_tree.zero_() tree_tree.zero_()
wp.launch( wp.launch(
@@ -184,14 +184,14 @@ def _flood_fill(
# Model: # Model:
ntree: int, ntree: int,
# In: # In:
tree_tree_in: wp.array3d(dtype=int), tree_tree_in: wp.array3d[int],
labels_in: wp.array2d(dtype=int), labels_in: wp.array2d[int],
stack_in: wp.array2d(dtype=int), stack_in: wp.array2d[int],
# Data out: # Data out:
nisland_out: wp.array(dtype=int), nisland_out: wp.array[int],
tree_island_out: wp.array2d(dtype=int), tree_island_out: wp.array2d[int],
# Out: # Out:
stack_out: wp.array2d(dtype=int), stack_out: wp.array2d[int],
): ):
"""DFS flood fill to discover islands using tree_tree matrix.""" """DFS flood fill to discover islands using tree_tree matrix."""
worldid = wp.tid() worldid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+87 -87
View File
@@ -71,18 +71,18 @@ def _ellipsoid_max_moment(size: wp.vec3, dir: int) -> float:
def _spring_damper_dof_passive( def _spring_damper_dof_passive(
# Model: # Model:
opt_disableflags: int, opt_disableflags: int,
qpos_spring: wp.array2d(dtype=float), qpos_spring: wp.array2d[float],
jnt_type: wp.array(dtype=int), jnt_type: wp.array[int],
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
jnt_stiffness: wp.array2d(dtype=float), jnt_stiffness: wp.array2d[float],
dof_damping: wp.array2d(dtype=float), dof_damping: wp.array2d[float],
# Data in: # Data in:
qpos_in: wp.array2d(dtype=float), qpos_in: wp.array2d[float],
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
# Data out: # Data out:
qfrc_spring_out: wp.array2d(dtype=float), qfrc_spring_out: wp.array2d[float],
qfrc_damper_out: wp.array2d(dtype=float), qfrc_damper_out: wp.array2d[float],
): ):
worldid, jntid = wp.tid() worldid, jntid = wp.tid()
dofid = jnt_dofadr[jntid] dofid = jnt_dofadr[jntid]
@@ -182,22 +182,22 @@ def _spring_damper_dof_passive(
@wp.kernel @wp.kernel
def _spring_damper_tendon_passive( def _spring_damper_tendon_passive(
# Model: # Model:
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
tendon_stiffness: wp.array2d(dtype=float), tendon_stiffness: wp.array2d[float],
tendon_damping: wp.array2d(dtype=float), tendon_damping: wp.array2d[float],
tendon_lengthspring: wp.array2d(dtype=wp.vec2), tendon_lengthspring: wp.array2d[wp.vec2],
# Data in: # Data in:
ten_J_in: wp.array2d(dtype=float), ten_J_in: wp.array2d[float],
ten_length_in: wp.array2d(dtype=float), ten_length_in: wp.array2d[float],
ten_velocity_in: wp.array2d(dtype=float), ten_velocity_in: wp.array2d[float],
# In: # In:
dsbl_spring: bool, dsbl_spring: bool,
dsbl_damper: bool, dsbl_damper: bool,
# Data out: # Data out:
qfrc_spring_out: wp.array2d(dtype=float), qfrc_spring_out: wp.array2d[float],
qfrc_damper_out: wp.array2d(dtype=float), qfrc_damper_out: wp.array2d[float],
): ):
worldid, tenid, dofid_sparse = wp.tid() worldid, tenid, dofid_sparse = wp.tid()
@@ -246,18 +246,18 @@ def _spring_damper_tendon_passive(
@wp.kernel @wp.kernel
def _gravity_force( def _gravity_force(
# Model: # Model:
opt_gravity: wp.array(dtype=wp.vec3), opt_gravity: wp.array[wp.vec3],
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_mass: wp.array2d(dtype=float), body_mass: wp.array2d[float],
body_gravcomp: wp.array2d(dtype=float), body_gravcomp: wp.array2d[float],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
# Data in: # Data in:
xipos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
# Data out: # Data out:
qfrc_gravcomp_out: wp.array2d(dtype=float), qfrc_gravcomp_out: wp.array2d[float],
): ):
worldid, bodyid, dofid = wp.tid() worldid, bodyid, dofid = wp.tid()
bodyid += 1 # skip world body bodyid += 1 # skip world body
@@ -275,27 +275,27 @@ def _gravity_force(
@wp.kernel @wp.kernel
def _fluid_force( def _fluid_force(
# Model: # Model:
opt_wind: wp.array(dtype=wp.vec3), opt_wind: wp.array[wp.vec3],
opt_density: wp.array(dtype=float), opt_density: wp.array[float],
opt_viscosity: wp.array(dtype=float), opt_viscosity: wp.array[float],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_geomnum: wp.array(dtype=int), body_geomnum: wp.array[int],
body_geomadr: wp.array(dtype=int), body_geomadr: wp.array[int],
body_mass: wp.array2d(dtype=float), body_mass: wp.array2d[float],
body_inertia: wp.array2d(dtype=wp.vec3), body_inertia: wp.array2d[wp.vec3],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_fluid: wp.array2d(dtype=float), geom_fluid: wp.array2d[float],
body_fluid_ellipsoid: wp.array(dtype=bool), body_fluid_ellipsoid: wp.array[bool],
# Data in: # Data in:
xipos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d[wp.vec3],
ximat_in: wp.array2d(dtype=wp.mat33), ximat_in: wp.array2d[wp.mat33],
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cvel_in: wp.array2d(dtype=wp.spatial_vector), cvel_in: wp.array2d[wp.spatial_vector],
# Out: # Out:
fluid_applied_out: wp.array2d(dtype=wp.spatial_vector), fluid_applied_out: wp.array2d[wp.spatial_vector],
): ):
"""Computes body-space fluid forces for both inertia-box and ellipsoid models.""" """Computes body-space fluid forces for both inertia-box and ellipsoid models."""
worldid, bodyid = wp.tid() worldid, bodyid = wp.tid()
@@ -535,18 +535,18 @@ def _fluid(m: Model, d: Data):
@wp.kernel @wp.kernel
def _qfrc_passive( def _qfrc_passive(
# Model: # Model:
jnt_actgravcomp: wp.array(dtype=int), jnt_actgravcomp: wp.array[int],
dof_jntid: wp.array(dtype=int), dof_jntid: wp.array[int],
has_fluid: bool, has_fluid: bool,
# Data in: # Data in:
qfrc_spring_in: wp.array2d(dtype=float), qfrc_spring_in: wp.array2d[float],
qfrc_damper_in: wp.array2d(dtype=float), qfrc_damper_in: wp.array2d[float],
qfrc_gravcomp_in: wp.array2d(dtype=float), qfrc_gravcomp_in: wp.array2d[float],
qfrc_fluid_in: wp.array2d(dtype=float), qfrc_fluid_in: wp.array2d[float],
# In: # In:
gravcomp: bool, gravcomp: bool,
# Data out: # Data out:
qfrc_passive_out: wp.array2d(dtype=float), qfrc_passive_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
qfrc_passive = qfrc_spring_in[worldid, dofid] qfrc_passive = qfrc_spring_in[worldid, dofid]
@@ -567,29 +567,29 @@ def _qfrc_passive(
def _flex_elasticity( def _flex_elasticity(
# Model: # Model:
nflex: int, nflex: int,
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array[float],
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_edgeadr: wp.array(dtype=int), flex_edgeadr: wp.array[int],
flex_elemadr: wp.array(dtype=int), flex_elemadr: wp.array[int],
flex_elemnum: wp.array(dtype=int), flex_elemnum: wp.array[int],
flex_elemdataadr: wp.array(dtype=int), flex_elemdataadr: wp.array[int],
flex_elemedgeadr: wp.array(dtype=int), flex_elemedgeadr: wp.array[int],
flex_vertbodyid: wp.array(dtype=int), flex_vertbodyid: wp.array[int],
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
flex_elemedge: wp.array(dtype=int), flex_elemedge: wp.array[int],
flexedge_length0: wp.array(dtype=float), flexedge_length0: wp.array[float],
flex_stiffness: wp.array2d(dtype=float), flex_stiffness: wp.array2d[float],
flex_damping: wp.array(dtype=float), flex_damping: wp.array[float],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
flexedge_length_in: wp.array2d(dtype=float), flexedge_length_in: wp.array2d[float],
flexedge_velocity_in: wp.array2d(dtype=float), flexedge_velocity_in: wp.array2d[float],
# In: # In:
dsbl_damper: bool, dsbl_damper: bool,
# Data out: # Data out:
qfrc_spring_out: wp.array2d(dtype=float), qfrc_spring_out: wp.array2d[float],
): ):
worldid, elemid = wp.tid() worldid, elemid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]] timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -665,19 +665,19 @@ def _flex_elasticity(
def _flex_bending( def _flex_bending(
# Model: # Model:
nflex: int, nflex: int,
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_edgeadr: wp.array(dtype=int), flex_edgeadr: wp.array[int],
flex_edgenum: wp.array(dtype=int), flex_edgenum: wp.array[int],
flex_vertbodyid: wp.array(dtype=int), flex_vertbodyid: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_edgeflap: wp.array(dtype=wp.vec2i), flex_edgeflap: wp.array[wp.vec2i],
flex_bending: wp.array2d(dtype=float), flex_bending: wp.array2d[float],
# Data in: # Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# Data out: # Data out:
qfrc_spring_out: wp.array2d(dtype=float), qfrc_spring_out: wp.array2d[float],
): ):
worldid, edgeid = wp.tid() worldid, edgeid = wp.tid()
nvert = 4 nvert = 4
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+114 -112
View File
@@ -52,12 +52,12 @@ def _ray_map(pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w
@wp.func @wp.func
def _ray_eliminate( def _ray_eliminate(
# Model: # Model:
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_matid: wp.array(dtype=int), # kernel_analyzer: ignore geom_matid: wp.array[int], # kernel_analyzer: ignore
geom_group: wp.array(dtype=int), geom_group: wp.array[int],
geom_rgba: wp.array(dtype=wp.vec4), # kernel_analyzer: ignore geom_rgba: wp.array[wp.vec4], # kernel_analyzer: ignore
mat_rgba: wp.array(dtype=wp.vec4), # kernel_analyzer: ignore mat_rgba: wp.array[wp.vec4], # kernel_analyzer: ignore
# In: # In:
geomid: int, geomid: int,
geomgroup: vec6, geomgroup: vec6,
@@ -451,26 +451,27 @@ def ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.ve
@wp.func @wp.func
def ray_hfield( def ray_hfield(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
hfield_size: wp.array(dtype=wp.vec4), hfield_size: wp.array[wp.vec4],
hfield_nrow: wp.array(dtype=int), hfield_nrow: wp.array[int],
hfield_ncol: wp.array(dtype=int), hfield_ncol: wp.array[int],
hfield_adr: wp.array(dtype=int), hfield_adr: wp.array[int],
hfield_data: wp.array(dtype=float), hfield_data: wp.array[float],
# In: # In:
pos: wp.vec3, pos: wp.vec3,
mat: wp.mat33, mat: wp.mat33,
pnt: wp.vec3, pnt: wp.vec3,
vec: wp.vec3, vec: wp.vec3,
id: int, id: int,
worldid: int,
) -> Tuple[float, wp.vec3]: ) -> Tuple[float, wp.vec3]:
# check geom type # check geom type
if geom_type[id] != GeomType.HFIELD: if geom_type[id] != GeomType.HFIELD:
return -1.0, wp.vec3() return -1.0, wp.vec3()
# hfield id and dimensions # hfield id and dimensions
hid = geom_dataid[id] hid = geom_dataid[worldid % geom_dataid.shape[0], id]
nrow = hfield_nrow[hid] nrow = hfield_nrow[hid]
ncol = hfield_ncol[hid] ncol = hfield_ncol[hid]
@@ -622,10 +623,10 @@ def ray_hfield(
def ray_mesh( def ray_mesh(
# Model: # Model:
nmeshface: int, nmeshface: int,
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_face: wp.array(dtype=wp.vec3i), mesh_face: wp.array[wp.vec3i],
# In: # In:
data_id: int, data_id: int,
pos: wp.vec3, pos: wp.vec3,
@@ -699,7 +700,7 @@ def ray_mesh(
@wp.func @wp.func
def ray_mesh_with_bvh( def ray_mesh_with_bvh(
# In: # In:
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
mesh_geom_id: int, mesh_geom_id: int,
pos: wp.vec3, pos: wp.vec3,
mat: wp.mat33, mat: wp.mat33,
@@ -732,7 +733,7 @@ def ray_mesh_with_bvh(
@wp.func @wp.func
def ray_mesh_with_bvh_anyhit( def ray_mesh_with_bvh_anyhit(
# In: # In:
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
mesh_geom_id: int, mesh_geom_id: int,
pos: wp.vec3, pos: wp.vec3,
mat: wp.mat33, mat: wp.mat33,
@@ -752,7 +753,7 @@ def ray_mesh_with_bvh_anyhit(
@wp.func @wp.func
def ray_flex_with_bvh( def ray_flex_with_bvh(
# In: # In:
flex_bvh_id: wp.array(dtype=wp.uint64), flex_bvh_id: wp.array[wp.uint64],
flexid: int, flexid: int,
group_root: int, group_root: int,
pnt: wp.vec3, pnt: wp.vec3,
@@ -781,7 +782,7 @@ def ray_flex_with_bvh(
@wp.func @wp.func
def ray_flex_with_bvh_anyhit( def ray_flex_with_bvh_anyhit(
# In: # In:
flex_bvh_id: wp.array(dtype=wp.uint64), flex_bvh_id: wp.array[wp.uint64],
flexid: int, flexid: int,
group_root: int, group_root: int,
pnt: wp.vec3, pnt: wp.vec3,
@@ -823,27 +824,27 @@ def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
def _ray_geom_mesh( def _ray_geom_mesh(
# Model: # Model:
nmeshface: int, nmeshface: int,
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_matid: wp.array2d(dtype=int), geom_matid: wp.array2d[int],
geom_group: wp.array(dtype=int), geom_group: wp.array[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rgba: wp.array2d(dtype=wp.vec4), geom_rgba: wp.array2d[wp.vec4],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_face: wp.array(dtype=wp.vec3i), mesh_face: wp.array[wp.vec3i],
hfield_size: wp.array(dtype=wp.vec4), hfield_size: wp.array[wp.vec4],
hfield_nrow: wp.array(dtype=int), hfield_nrow: wp.array[int],
hfield_ncol: wp.array(dtype=int), hfield_ncol: wp.array[int],
hfield_adr: wp.array(dtype=int), hfield_adr: wp.array[int],
hfield_data: wp.array(dtype=float), hfield_data: wp.array[float],
mat_rgba: wp.array2d(dtype=wp.vec4), mat_rgba: wp.array2d[wp.vec4],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
worldid: int, worldid: int,
pnt: wp.vec3, pnt: wp.vec3,
@@ -876,7 +877,7 @@ def _ray_geom_mesh(
mesh_faceadr, mesh_faceadr,
mesh_vert, mesh_vert,
mesh_face, mesh_face,
geom_dataid[geomid], geom_dataid[worldid % geom_dataid.shape[0], geomid],
pos, pos,
mat, mat,
geom_size[worldid % geom_size.shape[0], geomid], geom_size[worldid % geom_size.shape[0], geomid],
@@ -897,6 +898,7 @@ def _ray_geom_mesh(
pnt, pnt,
vec, vec,
geomid, geomid,
worldid,
) )
else: else:
return ray_geom(pos, mat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type) return ray_geom(pos, mat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type)
@@ -909,37 +911,37 @@ def _ray(
# Model: # Model:
ngeom: int, ngeom: int,
nmeshface: int, nmeshface: int,
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_matid: wp.array2d(dtype=int), geom_matid: wp.array2d[int],
geom_group: wp.array(dtype=int), geom_group: wp.array[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rgba: wp.array2d(dtype=wp.vec4), geom_rgba: wp.array2d[wp.vec4],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_face: wp.array(dtype=wp.vec3i), mesh_face: wp.array[wp.vec3i],
hfield_size: wp.array(dtype=wp.vec4), hfield_size: wp.array[wp.vec4],
hfield_nrow: wp.array(dtype=int), hfield_nrow: wp.array[int],
hfield_ncol: wp.array(dtype=int), hfield_ncol: wp.array[int],
hfield_adr: wp.array(dtype=int), hfield_adr: wp.array[int],
hfield_data: wp.array(dtype=float), hfield_data: wp.array[float],
mat_rgba: wp.array2d(dtype=wp.vec4), mat_rgba: wp.array2d[wp.vec4],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
pnt: wp.array2d(dtype=wp.vec3), pnt: wp.array2d[wp.vec3],
vec: wp.array2d(dtype=wp.vec3), vec: wp.array2d[wp.vec3],
geomgroup: vec6, geomgroup: vec6,
flg_static: bool, flg_static: bool,
bodyexclude: wp.array(dtype=int), bodyexclude: wp.array[int],
# Out: # Out:
dist_out: wp.array2d(dtype=float), dist_out: wp.array2d[float],
geomid_out: wp.array2d(dtype=int), geomid_out: wp.array2d[int],
normal_out: wp.array2d(dtype=wp.vec3), normal_out: wp.array2d[wp.vec3],
): ):
worldid, rayid, tid = wp.tid() worldid, rayid, tid = wp.tid()
@@ -1011,18 +1013,18 @@ def _ray(
@wp.func @wp.func
def _ray_geom_mesh_bvh( def _ray_geom_mesh_bvh(
# Model: # Model:
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_matid: wp.array2d(dtype=int), geom_matid: wp.array2d[int],
geom_group: wp.array(dtype=int), geom_group: wp.array[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rgba: wp.array2d(dtype=wp.vec4), geom_rgba: wp.array2d[wp.vec4],
mat_rgba: wp.array2d(dtype=wp.vec4), mat_rgba: wp.array2d[wp.vec4],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
worldid: int, worldid: int,
pnt: wp.vec3, pnt: wp.vec3,
@@ -1031,8 +1033,8 @@ def _ray_geom_mesh_bvh(
flg_static: bool, flg_static: bool,
bodyexclude: int, bodyexclude: int,
geomid: int, geomid: int,
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array(dtype=wp.uint64), hfield_bvh_id: wp.array[wp.uint64],
min_dist: float, min_dist: float,
) -> Tuple[float, wp.vec3]: ) -> Tuple[float, wp.vec3]:
if not _ray_eliminate( if not _ray_eliminate(
@@ -1055,7 +1057,7 @@ def _ray_geom_mesh_bvh(
bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh( t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
bvh_ids, bvh_ids,
geom_dataid[geomid], geom_dataid[worldid % geom_dataid.shape[0], geomid],
pos, pos,
mat, mat,
pnt, pnt,
@@ -1081,33 +1083,33 @@ def _ray_geom_mesh_bvh(
def _ray_bvh( def _ray_bvh(
# Model: # Model:
ngeom: int, ngeom: int,
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_matid: wp.array2d(dtype=int), geom_matid: wp.array2d[int],
geom_group: wp.array(dtype=int), geom_group: wp.array[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rgba: wp.array2d(dtype=wp.vec4), geom_rgba: wp.array2d[wp.vec4],
mat_rgba: wp.array2d(dtype=wp.vec4), mat_rgba: wp.array2d[wp.vec4],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
# In: # In:
pnt: wp.array2d(dtype=wp.vec3), pnt: wp.array2d[wp.vec3],
vec: wp.array2d(dtype=wp.vec3), vec: wp.array2d[wp.vec3],
geomgroup: vec6, geomgroup: vec6,
flg_static: bool, flg_static: bool,
bodyexclude: wp.array(dtype=int), bodyexclude: wp.array[int],
bvh_id: wp.uint64, bvh_id: wp.uint64,
group_root: wp.array(dtype=int), group_root: wp.array[int],
enabled_geom_ids: wp.array(dtype=int), enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array(dtype=wp.uint64), hfield_bvh_id: wp.array[wp.uint64],
# Out: # Out:
dist_out: wp.array2d(dtype=float), dist_out: wp.array2d[float],
geomid_out: wp.array2d(dtype=int), geomid_out: wp.array2d[int],
normal_out: wp.array2d(dtype=wp.vec3), normal_out: wp.array2d[wp.vec3],
): ):
worldid, rayid = wp.tid() worldid, rayid = wp.tid()
@@ -1166,8 +1168,8 @@ def _ray_bvh(
def ray( def ray(
m: Model, m: Model,
d: Data, d: Data,
pnt: wp.array2d(dtype=wp.vec3), pnt: wp.array2d[wp.vec3],
vec: wp.array2d(dtype=wp.vec3), vec: wp.array2d[wp.vec3],
geomgroup: vec6 | None = None, geomgroup: vec6 | None = None,
flg_static: bool = True, flg_static: bool = True,
bodyexclude: int = -1, bodyexclude: int = -1,
@@ -1210,14 +1212,14 @@ def ray(
def rays( def rays(
m: Model, m: Model,
d: Data, d: Data,
pnt: wp.array2d(dtype=wp.vec3), pnt: wp.array2d[wp.vec3],
vec: wp.array2d(dtype=wp.vec3), vec: wp.array2d[wp.vec3],
geomgroup: vec6, geomgroup: vec6,
flg_static: bool, flg_static: bool,
bodyexclude: wp.array(dtype=int), bodyexclude: wp.array[int],
dist: wp.array2d(dtype=float), dist: wp.array2d[float],
geomid: wp.array2d(dtype=int), geomid: wp.array2d[int],
normal: wp.array2d(dtype=wp.vec3), normal: wp.array2d[wp.vec3],
rc: RenderContext | None = None, rc: RenderContext | None = None,
): ):
"""Ray intersection for multiple worlds and multiple rays. """Ray intersection for multiple worlds and multiple rays.
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+108 -108
View File
@@ -43,17 +43,17 @@ wp.set_module_options({"enable_backward": False})
@wp.func @wp.func
def sample_texture( def sample_texture(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
# In: # In:
geom_id: int, geom_id: int,
tex_repeat: wp.vec2, tex_repeat: wp.vec2,
tex: wp.Texture2D, tex: wp.Texture2D,
pos: wp.vec3, pos: wp.vec3,
rot: wp.mat33, rot: wp.mat33,
mesh_facetexcoord: wp.array(dtype=wp.vec3i), mesh_facetexcoord: wp.array[wp.vec3i],
mesh_texcoord: wp.array(dtype=wp.vec2), mesh_texcoord: wp.array[wp.vec2],
mesh_texcoord_offsets: wp.array(dtype=int), mesh_texcoord_offsets: wp.array[int],
hit_point: wp.vec3, hit_point: wp.vec3,
bary_u: float, bary_u: float,
bary_v: float, bary_v: float,
@@ -88,29 +88,29 @@ def sample_texture(
@wp.func @wp.func
def cast_ray( def cast_ray(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
bvh_id: wp.uint64, bvh_id: wp.uint64,
group_root: int, group_root: int,
worldid: int, worldid: int,
bvh_ngeom: int, bvh_ngeom: int,
flex_bvh_ngeom: int, flex_bvh_ngeom: int,
enabled_geom_ids: wp.array(dtype=int), enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array(dtype=wp.uint64), hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array(dtype=int), flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array(dtype=int), flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array(dtype=wp.uint64), flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d(dtype=int), flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3, ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3, ray_dir_world: wp.vec3,
) -> Tuple[int, float, wp.vec3, float, float, int, int]: ) -> Tuple[int, float, wp.vec3, float, float, int, int]:
@@ -159,7 +159,7 @@ def cast_ray(
if gtype == GeomType.HFIELD: if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id, hfield_bvh_id,
geom_dataid[gi], geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi], geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi], geom_xmat_in[worldid, gi],
ray_origin_world, ray_origin_world,
@@ -208,7 +208,7 @@ def cast_ray(
if gtype == GeomType.MESH: if gtype == GeomType.MESH:
d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh( d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
mesh_bvh_id, mesh_bvh_id,
geom_dataid[gi], geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi], geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi], geom_xmat_in[worldid, gi],
ray_origin_world, ray_origin_world,
@@ -256,29 +256,29 @@ def cast_ray(
@wp.func @wp.func
def cast_ray_first_hit( def cast_ray_first_hit(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
bvh_id: wp.uint64, bvh_id: wp.uint64,
group_root: int, group_root: int,
worldid: int, worldid: int,
bvh_ngeom: int, bvh_ngeom: int,
bvh_nflexgeom: int, bvh_nflexgeom: int,
enabled_geom_ids: wp.array(dtype=int), enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array(dtype=wp.uint64), hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array(dtype=int), flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array(dtype=int), flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array(dtype=wp.uint64), flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d(dtype=int), flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3, ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3, ray_dir_world: wp.vec3,
max_dist: float, max_dist: float,
@@ -314,7 +314,7 @@ def cast_ray_first_hit(
if gtype == GeomType.HFIELD: if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id, hfield_bvh_id,
geom_dataid[gi], geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi], geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi], geom_xmat_in[worldid, gi],
ray_origin_world, ray_origin_world,
@@ -363,7 +363,7 @@ def cast_ray_first_hit(
if gtype == GeomType.MESH: if gtype == GeomType.MESH:
hit = ray_mesh_with_bvh_anyhit( hit = ray_mesh_with_bvh_anyhit(
mesh_bvh_id, mesh_bvh_id,
geom_dataid[gi], geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi], geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi], geom_xmat_in[worldid, gi],
ray_origin_world, ray_origin_world,
@@ -409,30 +409,30 @@ def cast_ray_first_hit(
@wp.func @wp.func
def compute_lighting( def compute_lighting(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
use_shadows: bool, use_shadows: bool,
bvh_id: wp.uint64, bvh_id: wp.uint64,
group_root: int, group_root: int,
bvh_ngeom: int, bvh_ngeom: int,
bvh_nflexgeom: int, bvh_nflexgeom: int,
enabled_geom_ids: wp.array(dtype=int), enabled_geom_ids: wp.array[int],
worldid: int, worldid: int,
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array(dtype=wp.uint64), hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array(dtype=int), flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array(dtype=int), flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array(dtype=wp.uint64), flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d(dtype=int), flex_group_root: wp.array2d[int],
lightactive: bool, lightactive: bool,
lighttype: int, lighttype: int,
lightcastshadow: bool, lightcastshadow: bool,
@@ -530,65 +530,65 @@ def render(m: Model, d: Data, rc: RenderContext):
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def _render_megakernel( def _render_megakernel(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_matid: wp.array2d(dtype=int), geom_matid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_rgba: wp.array2d(dtype=wp.vec4), geom_rgba: wp.array2d[wp.vec4],
cam_projection: wp.array(dtype=int), cam_projection: wp.array[int],
cam_fovy: wp.array2d(dtype=float), cam_fovy: wp.array2d[float],
cam_sensorsize: wp.array(dtype=wp.vec2), cam_sensorsize: wp.array[wp.vec2],
cam_intrinsic: wp.array2d(dtype=wp.vec4), cam_intrinsic: wp.array2d[wp.vec4],
light_type: wp.array2d(dtype=int), light_type: wp.array2d[int],
light_castshadow: wp.array2d(dtype=bool), light_castshadow: wp.array2d[bool],
light_active: wp.array2d(dtype=bool), light_active: wp.array2d[bool],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
mat_texid: wp.array3d(dtype=int), mat_texid: wp.array3d[int],
mat_texrepeat: wp.array2d(dtype=wp.vec2), mat_texrepeat: wp.array2d[wp.vec2],
mat_rgba: wp.array2d(dtype=wp.vec4), mat_rgba: wp.array2d[wp.vec4],
# Data in: # Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d[wp.mat33],
cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xpos_in: wp.array2d[wp.vec3],
cam_xmat_in: wp.array2d(dtype=wp.mat33), cam_xmat_in: wp.array2d[wp.mat33],
light_xpos_in: wp.array2d(dtype=wp.vec3), light_xpos_in: wp.array2d[wp.vec3],
light_xdir_in: wp.array2d(dtype=wp.vec3), light_xdir_in: wp.array2d[wp.vec3],
flexvert_xpos_in: wp.array2d(dtype=wp.vec3), flexvert_xpos_in: wp.array2d[wp.vec3],
# In: # In:
nrender: int, nrender: int,
use_shadows: bool, use_shadows: bool,
bvh_ngeom: int, bvh_ngeom: int,
bvh_nflexgeom: int, bvh_nflexgeom: int,
cam_res: wp.array(dtype=wp.vec2i), cam_res: wp.array[wp.vec2i],
cam_id_map: wp.array(dtype=int), cam_id_map: wp.array[int],
ray: wp.array(dtype=wp.vec3), ray: wp.array[wp.vec3],
rgb_adr: wp.array(dtype=int), rgb_adr: wp.array[int],
depth_adr: wp.array(dtype=int), depth_adr: wp.array[int],
seg_adr: wp.array(dtype=int), seg_adr: wp.array[int],
render_rgb: wp.array(dtype=bool), render_rgb: wp.array[bool],
render_depth: wp.array(dtype=bool), render_depth: wp.array[bool],
render_seg: wp.array(dtype=bool), render_seg: wp.array[bool],
bvh_id: wp.uint64, bvh_id: wp.uint64,
group_root: wp.array(dtype=int), group_root: wp.array[int],
flex_bvh_id: wp.array(dtype=wp.uint64), flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d(dtype=int), flex_group_root: wp.array2d[int],
enabled_geom_ids: wp.array(dtype=int), enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array(dtype=wp.uint64), mesh_bvh_id: wp.array[wp.uint64],
mesh_facetexcoord: wp.array(dtype=wp.vec3i), mesh_facetexcoord: wp.array[wp.vec3i],
mesh_texcoord: wp.array(dtype=wp.vec2), mesh_texcoord: wp.array[wp.vec2],
mesh_texcoord_offsets: wp.array(dtype=int), mesh_texcoord_offsets: wp.array[int],
hfield_bvh_id: wp.array(dtype=wp.uint64), hfield_bvh_id: wp.array[wp.uint64],
flex_rgba: wp.array(dtype=wp.vec4), flex_rgba: wp.array[wp.vec4],
flex_geom_flexid: wp.array(dtype=int), flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array(dtype=int), flex_geom_edgeid: wp.array[int],
textures: wp.array(dtype=wp.Texture2D), textures: wp.array[wp.Texture2D],
# Out: # Out:
rgb_out: wp.array2d(dtype=wp.uint32), rgb_out: wp.array2d[wp.uint32],
depth_out: wp.array2d(dtype=float), depth_out: wp.array2d[float],
seg_out: wp.array2d(dtype=int), seg_out: wp.array2d[int],
): ):
worldid, rayid = wp.tid() worldid, rayid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
Vendored
+14 -14
View File
@@ -28,9 +28,9 @@ def _convert_texture_data(
width: int, width: int,
adr: int, adr: int,
nc: int, nc: int,
tex_data_in: wp.array(dtype=wp.uint8), tex_data_in: wp.array[wp.uint8],
# Out: # Out:
tex_data_out: wp.array3d(dtype=float), tex_data_out: wp.array3d[float],
): ):
"""Convert uint8 texture data to vec4 format for efficient sampling.""" """Convert uint8 texture data to vec4 format for efficient sampling."""
x, y = wp.tid() x, y = wp.tid()
@@ -134,11 +134,11 @@ def pack_rgba_to_uint32(r: float, g: float, b: float, a: float) -> wp.uint32:
@wp.kernel @wp.kernel
def unpack_rgb_kernel( def unpack_rgb_kernel(
# In: # In:
packed: wp.array2d(dtype=wp.uint32), packed: wp.array2d[wp.uint32],
rgb_adr: wp.array(dtype=int), rgb_adr: wp.array[int],
camera_index: int, camera_index: int,
# Out: # Out:
rgb_out: wp.array3d(dtype=wp.vec3), rgb_out: wp.array3d[wp.vec3],
): ):
"""Unpack ABGR uint32 packed pixel data into separate R, G, and B channels.""" """Unpack ABGR uint32 packed pixel data into separate R, G, and B channels."""
worldid, pixelid = wp.tid() worldid, pixelid = wp.tid()
@@ -157,12 +157,12 @@ def unpack_rgb_kernel(
@wp.kernel @wp.kernel
def extract_depth_kernel( def extract_depth_kernel(
# In: # In:
depth_data: wp.array2d(dtype=float), depth_data: wp.array2d[float],
depth_adr: wp.array(dtype=int), depth_adr: wp.array[int],
camera_index: int, camera_index: int,
depth_scale: float, depth_scale: float,
# Out: # Out:
depth_out: wp.array3d(dtype=float), depth_out: wp.array3d[float],
): ):
"""Extract the depth data from the render context buffers for a given camera index.""" """Extract the depth data from the render context buffers for a given camera index."""
worldid, pixelid = wp.tid() worldid, pixelid = wp.tid()
@@ -174,7 +174,7 @@ def extract_depth_kernel(
depth_out[worldid, yid, xid] = wp.clamp(val / depth_scale, 0.0, 1.0) depth_out[worldid, yid, xid] = wp.clamp(val / depth_scale, 0.0, 1.0)
def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d(dtype=wp.vec3)): def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d[wp.vec3]):
"""Get the RGB data output from the render context buffers for a given camera index. """Get the RGB data output from the render context buffers for a given camera index.
Args: Args:
@@ -190,7 +190,7 @@ def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d(dtype=wp.v
) )
def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_out: wp.array3d(dtype=float)): def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_out: wp.array3d[float]):
"""Get the depth data output from the render context buffers for a given camera index. """Get the depth data output from the render context buffers for a given camera index.
Args: Args:
@@ -211,11 +211,11 @@ def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_ou
@wp.kernel @wp.kernel
def _extract_seg_kernel( def _extract_seg_kernel(
# In: # In:
seg_data: wp.array2d(dtype=int), seg_data: wp.array2d[int],
seg_adr: wp.array(dtype=int), seg_adr: wp.array[int],
camera_index: int, camera_index: int,
# Out: # Out:
seg_out: wp.array3d(dtype=int), seg_out: wp.array3d[int],
): ):
"""Extract per-pixel geom IDs from the render context buffers for a given camera index.""" """Extract per-pixel geom IDs from the render context buffers for a given camera index."""
worldid, pixelid = wp.tid() worldid, pixelid = wp.tid()
@@ -226,7 +226,7 @@ def _extract_seg_kernel(
seg_out[worldid, yid, xid] = seg_data[worldid, seg_adr_offset + pixelid] 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)): def get_segmentation(rc: RenderContext, camera_index: int, seg_out: wp.array3d[int]):
"""Get the segmentation data from the render context buffers for a given camera index. """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 Each pixel contains the MuJoCo geom ID of the geometry hit by the ray, -1 for
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+442 -442
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+471 -471
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+324 -324
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+96 -100
View File
@@ -69,16 +69,16 @@ def mul_m_sparse(check_skip: bool):
@wp.kernel(module="unique") @wp.kernel(module="unique")
def _mul_m_sparse( def _mul_m_sparse(
# Model: # Model:
qM_mulm_rowadr: wp.array(dtype=int), qM_mulm_rowadr: wp.array[int],
qM_mulm_col: wp.array(dtype=int), qM_mulm_col: wp.array[int],
qM_mulm_madr: wp.array(dtype=int), qM_mulm_madr: wp.array[int],
# Data in: # Data in:
qM_in: wp.array3d(dtype=float), qM_in: wp.array3d[float],
# In: # In:
vec: wp.array2d(dtype=float), vec: wp.array2d[float],
skip: wp.array(dtype=bool), skip: wp.array[bool],
# Out: # Out:
res: wp.array2d(dtype=float), res: wp.array2d[float],
): ):
"""Sparse matmul: one thread per DOF, gather-based (no atomics).""" """Sparse matmul: one thread per DOF, gather-based (no atomics)."""
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
@@ -108,12 +108,12 @@ def mul_m_dense(nv: int, check_skip: bool):
@wp.kernel(module="unique") @wp.kernel(module="unique")
def _mul_m_dense( def _mul_m_dense(
# Data in: # Data in:
qM_in: wp.array3d(dtype=float), qM_in: wp.array3d[float],
# In: # In:
vec: wp.array2d(dtype=float), vec: wp.array2d[float],
skip: wp.array(dtype=bool), skip: wp.array[bool],
# Out: # Out:
res: wp.array2d(dtype=float), res: wp.array2d[float],
): ):
worldid, i = wp.tid() worldid, i = wp.tid()
@@ -133,8 +133,8 @@ def mul_m_dense(nv: int, check_skip: bool):
def mul_m( def mul_m(
m: Model, m: Model,
d: Data, d: Data,
res: wp.array2d(dtype=float), res: wp.array2d[float],
vec: wp.array2d(dtype=float), vec: wp.array2d[float],
skip: Optional[wp.array] = None, skip: Optional[wp.array] = None,
M: Optional[wp.array] = None, M: Optional[wp.array] = None,
): ):
@@ -175,18 +175,18 @@ def mul_m(
def _apply_ft( def _apply_ft(
# Model: # Model:
nbody: int, nbody: int,
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
# Data in: # Data in:
xipos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
# In: # In:
ft_in: wp.array2d(dtype=wp.spatial_vector), ft_in: wp.array2d[wp.spatial_vector],
flg_add: bool, flg_add: bool,
# Out: # Out:
qfrc_out: wp.array2d(dtype=float), qfrc_out: wp.array2d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
cdof = cdof_in[worldid, dofid] cdof = cdof_in[worldid, dofid]
@@ -216,7 +216,7 @@ def _apply_ft(
qfrc_out[worldid, dofid] = accumul qfrc_out[worldid, dofid] = accumul
def apply_ft(m: Model, d: Data, ft: wp.array2d(dtype=wp.spatial_vector), qfrc: wp.array2d(dtype=float), flg_add: bool): def apply_ft(m: Model, d: Data, ft: wp.array2d[wp.spatial_vector], qfrc: wp.array2d[float], flg_add: bool):
wp.launch( wp.launch(
kernel=_apply_ft, kernel=_apply_ft,
dim=(d.nworld, m.nv), dim=(d.nworld, m.nv),
@@ -226,7 +226,7 @@ def apply_ft(m: Model, d: Data, ft: wp.array2d(dtype=wp.spatial_vector), qfrc: w
@event_scope @event_scope
def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d[float]):
"""Map applied forces at each body via Jacobians to dof space and accumulate. """Map applied forces at each body via Jacobians to dof space and accumulate.
Args: Args:
@@ -238,9 +238,7 @@ def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)):
@wp.func @wp.func
def _decode_pyramid( def _decode_pyramid(njmax_in: int, pyramid: wp.array[float], efc_address: int, mu: vec5, condim: int) -> wp.spatial_vector:
njmax_in: int, pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int
) -> wp.spatial_vector:
"""Converts pyramid representation to contact force.""" """Converts pyramid representation to contact force."""
force = wp.spatial_vector() force = wp.spatial_vector()
@@ -270,13 +268,13 @@ def contact_force_fn(
# Model: # Model:
opt_cone: int, opt_cone: int,
# Data in: # Data in:
contact_frame_in: wp.array(dtype=wp.mat33), contact_frame_in: wp.array[wp.mat33],
contact_friction_in: wp.array(dtype=vec5), contact_friction_in: wp.array[vec5],
contact_dim_in: wp.array(dtype=int), contact_dim_in: wp.array[int],
contact_efc_address_in: wp.array2d(dtype=int), contact_efc_address_in: wp.array2d[int],
efc_force_in: wp.array2d(dtype=float), efc_force_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
nacon_in: wp.array(dtype=int), nacon_in: wp.array[int],
# In: # In:
worldid: int, worldid: int,
contact_id: int, contact_id: int,
@@ -315,19 +313,19 @@ def contact_force_kernel(
# Model: # Model:
opt_cone: int, opt_cone: int,
# Data in: # Data in:
contact_frame_in: wp.array(dtype=wp.mat33), contact_frame_in: wp.array[wp.mat33],
contact_friction_in: wp.array(dtype=vec5), contact_friction_in: wp.array[vec5],
contact_dim_in: wp.array(dtype=int), contact_dim_in: wp.array[int],
contact_efc_address_in: wp.array2d(dtype=int), contact_efc_address_in: wp.array2d[int],
contact_worldid_in: wp.array(dtype=int), contact_worldid_in: wp.array[int],
efc_force_in: wp.array2d(dtype=float), efc_force_in: wp.array2d[float],
njmax_in: int, njmax_in: int,
nacon_in: wp.array(dtype=int), nacon_in: wp.array[int],
# In: # In:
contact_ids: wp.array(dtype=int), contact_ids: wp.array[int],
to_world_frame: bool, to_world_frame: bool,
# Out: # Out:
out: wp.array(dtype=wp.spatial_vector), out: wp.array[wp.spatial_vector],
): ):
tid = wp.tid() tid = wp.tid()
@@ -353,9 +351,7 @@ def contact_force_kernel(
) )
def contact_force( def contact_force(m: Model, d: Data, contact_ids: wp.array[int], to_world_frame: bool, force: wp.array[wp.spatial_vector]):
m: Model, d: Data, contact_ids: wp.array(dtype=int), to_world_frame: bool, force: wp.array(dtype=wp.spatial_vector)
):
"""Compute forces for contacts in Data. """Compute forces for contacts in Data.
Args: Args:
@@ -400,12 +396,12 @@ def transform_force(frc: wp.spatial_vector, offset: wp.vec3) -> wp.spatial_vecto
@wp.func @wp.func
def jac_dof( def jac_dof(
# Model: # Model:
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
# Data in: # Data in:
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
# In: # In:
point: wp.vec3, point: wp.vec3,
bodyid: int, bodyid: int,
@@ -441,18 +437,18 @@ def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
@wp.kernel(module="unique", enable_backward=False) @wp.kernel(module="unique", enable_backward=False)
def _jac( def _jac(
# Model: # Model:
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
# Data in: # Data in:
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
# In: # In:
point_in: wp.array(dtype=wp.vec3), point_in: wp.array[wp.vec3],
bodyid_in: wp.array(dtype=int), bodyid_in: wp.array[int],
# Out: # Out:
jacp_out: wp.array3d(dtype=float), jacp_out: wp.array3d[float],
jacr_out: wp.array3d(dtype=float), jacr_out: wp.array3d[float],
): ):
worldid, dofid = wp.tid() worldid, dofid = wp.tid()
@@ -477,10 +473,10 @@ def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
def jac( def jac(
m: Model, m: Model,
d: Data, d: Data,
jacp: wp.array | None, # wp.array3d(dtype=float) jacp: wp.array | None, # wp.array3d[float]
jacr: wp.array | None, # wp.array3d(dtype=float) jacr: wp.array | None, # wp.array3d[float]
point: wp.array(dtype=wp.vec3), point: wp.array[wp.vec3],
body: wp.array(dtype=int), body: wp.array[int],
): ):
"""Compute translational and rotational Jacobian for point on body. """Compute translational and rotational Jacobian for point on body.
@@ -508,17 +504,17 @@ def jac(
@wp.func @wp.func
def jac_dot_dof( def jac_dot_dof(
# Model: # Model:
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
jnt_type: wp.array(dtype=int), jnt_type: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
dof_bodyid: wp.array(dtype=int), dof_bodyid: wp.array[int],
dof_jntid: wp.array(dtype=int), dof_jntid: wp.array[int],
# Data in: # Data in:
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d[wp.spatial_vector],
cvel_in: wp.array2d(dtype=wp.spatial_vector), cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d(dtype=wp.spatial_vector), cdof_dot_in: wp.array2d[wp.spatial_vector],
# In: # In:
point: wp.vec3, point: wp.vec3,
bodyid: int, bodyid: int,
@@ -573,7 +569,7 @@ def jac_dot_dof(
return jacp, jacr return jacp, jacr
def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, active: Optional[wp.array] = None): def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Optional[wp.array] = None):
"""Copy concatenated state components specified by sig from Data into state. """Copy concatenated state components specified by sig from Data into state.
The bits of the integer sig correspond to element fields of State. The bits of the integer sig correspond to element fields of State.
@@ -599,22 +595,22 @@ def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
neq: int, neq: int,
nmocap: int, nmocap: int,
# Data in: # Data in:
time_in: wp.array(dtype=float), time_in: wp.array[float],
qpos_in: wp.array2d(dtype=float), qpos_in: wp.array2d[float],
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d[float],
act_in: wp.array2d(dtype=float), act_in: wp.array2d[float],
qacc_warmstart_in: wp.array2d(dtype=float), qacc_warmstart_in: wp.array2d[float],
ctrl_in: wp.array2d(dtype=float), ctrl_in: wp.array2d[float],
qfrc_applied_in: wp.array2d(dtype=float), qfrc_applied_in: wp.array2d[float],
xfrc_applied_in: wp.array2d(dtype=wp.spatial_vector), xfrc_applied_in: wp.array2d[wp.spatial_vector],
eq_active_in: wp.array2d(dtype=bool), eq_active_in: wp.array2d[bool],
mocap_pos_in: wp.array2d(dtype=wp.vec3), mocap_pos_in: wp.array2d[wp.vec3],
mocap_quat_in: wp.array2d(dtype=wp.quat), mocap_quat_in: wp.array2d[wp.quat],
# In: # In:
sig_in: int, sig_in: int,
active_in: wp.array(dtype=bool), active_in: wp.array[bool],
# Out: # Out:
state_out: wp.array2d(dtype=float), state_out: wp.array2d[float],
): ):
worldid = wp.tid() worldid = wp.tid()
@@ -712,7 +708,7 @@ def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
) )
def set_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, active: Optional[wp.array] = None): def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Optional[wp.array] = None):
"""Copy concatenated state components specified by sig from state into Data. """Copy concatenated state components specified by sig from state into Data.
The bits of the integer sig correspond to element fields of State. The bits of the integer sig correspond to element fields of State.
@@ -739,20 +735,20 @@ def set_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
nmocap: int, nmocap: int,
# In: # In:
sig_in: int, sig_in: int,
active_in: wp.array(dtype=bool), active_in: wp.array[bool],
state_in: wp.array2d(dtype=float), state_in: wp.array2d[float],
# Data out: # Data out:
time_out: wp.array(dtype=float), time_out: wp.array[float],
qpos_out: wp.array2d(dtype=float), qpos_out: wp.array2d[float],
qvel_out: wp.array2d(dtype=float), qvel_out: wp.array2d[float],
act_out: wp.array2d(dtype=float), act_out: wp.array2d[float],
qacc_warmstart_out: wp.array2d(dtype=float), qacc_warmstart_out: wp.array2d[float],
ctrl_out: wp.array2d(dtype=float), ctrl_out: wp.array2d[float],
qfrc_applied_out: wp.array2d(dtype=float), qfrc_applied_out: wp.array2d[float],
xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector), xfrc_applied_out: wp.array2d[wp.spatial_vector],
eq_active_out: wp.array2d(dtype=bool), eq_active_out: wp.array2d[bool],
mocap_pos_out: wp.array2d(dtype=wp.vec3), mocap_pos_out: wp.array2d[wp.vec3],
mocap_quat_out: wp.array2d(dtype=wp.quat), mocap_quat_out: wp.array2d[wp.quat],
): ):
worldid = wp.tid() worldid = wp.tid()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+66 -56
View File
@@ -17,6 +17,7 @@ import enum
from typing import Callable from typing import Callable
import mujoco import mujoco
import numpy as np
import warp as wp import warp as wp
MJ_MINVAL = mujoco.mjMINVAL MJ_MINVAL = mujoco.mjMINVAL
@@ -793,9 +794,18 @@ class TileSet:
size: size of all the tiles in this set size: size of all the tiles in this set
""" """
adr: wp.array(dtype=int) adr: wp.array[int]
size: int size: int
def __eq__(self, other) -> bool:
if self.__class__ is not other.__class__:
return NotImplemented
return self.size == other.size and np.array_equal(np.asarray(self.adr.numpy()), np.asarray(other.adr.numpy()))
def __hash__(self) -> int:
adr = np.asarray(self.adr.numpy())
return hash((self.size, adr.dtype.str, adr.shape, adr.tobytes()))
@dataclasses.dataclass @dataclasses.dataclass
class Callback: class Callback:
@@ -933,7 +943,7 @@ class Model:
geom_conaffinity: geom contact affinity (ngeom,) geom_conaffinity: geom contact affinity (ngeom,)
geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,)
geom_bodyid: id of geom's body (ngeom,) geom_bodyid: id of geom's body (ngeom,)
geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) geom_dataid: id of geom's mesh/hfield; -1: none (*, ngeom)
geom_matid: material id for rendering (*, ngeom,) geom_matid: material id for rendering (*, ngeom,)
geom_group: geom group inclusion/exclusion mask (ngeom,) geom_group: geom group inclusion/exclusion mask (ngeom,)
geom_priority: geom contact priority (ngeom,) geom_priority: geom contact priority (ngeom,)
@@ -1322,7 +1332,7 @@ class Model:
geom_conaffinity: array("ngeom", int) geom_conaffinity: array("ngeom", int)
geom_condim: array("ngeom", int) geom_condim: array("ngeom", int)
geom_bodyid: array("ngeom", int) geom_bodyid: array("ngeom", int)
geom_dataid: array("ngeom", int) geom_dataid: array("*", "ngeom", int)
geom_matid: array("*", "ngeom", int) geom_matid: array("*", "ngeom", int)
geom_group: array("ngeom", int) geom_group: array("ngeom", int)
geom_priority: array("ngeom", int) geom_priority: array("ngeom", int)
@@ -1527,70 +1537,70 @@ class Model:
has_fluid: bool has_fluid: bool
has_sdf_geom: bool has_sdf_geom: bool
block_dim: BlockDim block_dim: BlockDim
body_tree: tuple[wp.array(dtype=int), ...] body_tree: tuple[wp.array[int], ...]
body_branches: wp.array(dtype=int) body_branches: wp.array[int]
body_branch_start: wp.array(dtype=int) body_branch_start: wp.array[int]
mocap_bodyid: array("nmocap", int) mocap_bodyid: array("nmocap", int)
body_fluid_ellipsoid: array("nbody", bool) body_fluid_ellipsoid: array("nbody", bool)
jnt_limited_slide_hinge_adr: wp.array(dtype=int) jnt_limited_slide_hinge_adr: wp.array[int]
jnt_limited_ball_adr: wp.array(dtype=int) jnt_limited_ball_adr: wp.array[int]
dof_tri_row: wp.array(dtype=int) dof_tri_row: wp.array[int]
dof_tri_col: wp.array(dtype=int) dof_tri_col: wp.array[int]
nxn_geom_pair: wp.array(dtype=wp.vec2i) nxn_geom_pair: wp.array[wp.vec2i]
nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i) nxn_geom_pair_filtered: wp.array[wp.vec2i]
nxn_pairid: wp.array(dtype=wp.vec2i) nxn_pairid: wp.array[wp.vec2i]
nxn_pairid_filtered: wp.array(dtype=wp.vec2i) nxn_pairid_filtered: wp.array[wp.vec2i]
geom_pair_type_count: tuple[int, ...] geom_pair_type_count: tuple[int, ...]
geom_plugin_index: array("ngeom", int) geom_plugin_index: array("ngeom", int)
eq_connect_adr: wp.array(dtype=int) eq_connect_adr: wp.array[int]
eq_wld_adr: wp.array(dtype=int) eq_wld_adr: wp.array[int]
eq_jnt_adr: wp.array(dtype=int) eq_jnt_adr: wp.array[int]
eq_ten_adr: wp.array(dtype=int) eq_ten_adr: wp.array[int]
eq_flex_adr: wp.array(dtype=int) eq_flex_adr: wp.array[int]
tendon_jnt_adr: wp.array(dtype=int) tendon_jnt_adr: wp.array[int]
tendon_site_pair_adr: wp.array(dtype=int) tendon_site_pair_adr: wp.array[int]
tendon_geom_adr: wp.array(dtype=int) tendon_geom_adr: wp.array[int]
tendon_limited_adr: wp.array(dtype=int) tendon_limited_adr: wp.array[int]
max_ten_J_rownnz: int max_ten_J_rownnz: int
ten_wrapadr_site: wp.array(dtype=int) ten_wrapadr_site: wp.array[int]
ten_wrapnum_site: wp.array(dtype=int) ten_wrapnum_site: wp.array[int]
wrap_jnt_adr: wp.array(dtype=int) wrap_jnt_adr: wp.array[int]
wrap_site_adr: wp.array(dtype=int) wrap_site_adr: wp.array[int]
wrap_site_pair_adr: wp.array(dtype=int) wrap_site_pair_adr: wp.array[int]
wrap_geom_adr: wp.array(dtype=int) wrap_geom_adr: wp.array[int]
wrap_pulley_scale: array("nwrap", float) wrap_pulley_scale: array("nwrap", float)
actuator_trntype_body_adr: wp.array(dtype=int) actuator_trntype_body_adr: wp.array[int]
sensor_pos_adr: wp.array(dtype=int) sensor_pos_adr: wp.array[int]
sensor_limitpos_adr: wp.array(dtype=int) sensor_limitpos_adr: wp.array[int]
sensor_vel_adr: wp.array(dtype=int) sensor_vel_adr: wp.array[int]
sensor_limitvel_adr: wp.array(dtype=int) sensor_limitvel_adr: wp.array[int]
sensor_acc_adr: wp.array(dtype=int) sensor_acc_adr: wp.array[int]
sensor_rangefinder_adr: wp.array(dtype=int) sensor_rangefinder_adr: wp.array[int]
rangefinder_sensor_adr: wp.array(dtype=int) rangefinder_sensor_adr: wp.array[int]
sensor_collision_start_adr: wp.array(dtype=int) sensor_collision_start_adr: wp.array[int]
collision_sensor_adr: array("nsensor", int) collision_sensor_adr: array("nsensor", int)
sensor_touch_adr: wp.array(dtype=int) sensor_touch_adr: wp.array[int]
sensor_limitfrc_adr: wp.array(dtype=int) sensor_limitfrc_adr: wp.array[int]
sensor_e_potential: bool sensor_e_potential: bool
sensor_e_kinetic: bool sensor_e_kinetic: bool
sensor_tendonactfrc_adr: wp.array(dtype=int) sensor_tendonactfrc_adr: wp.array[int]
sensor_subtree_vel: bool sensor_subtree_vel: bool
sensor_contact_adr: array("nsensorcontact", int) sensor_contact_adr: array("nsensorcontact", int)
sensor_adr_to_contact_adr: array("nsensor", int) sensor_adr_to_contact_adr: array("nsensor", int)
sensor_rne_postconstraint: bool sensor_rne_postconstraint: bool
sensor_rangefinder_bodyid: array("nrangefinder", int) sensor_rangefinder_bodyid: array("nrangefinder", int)
taxel_vertadr: array("nsensortaxel", int) taxel_vertadr: array("nsensortaxel", int)
taxel_sensorid: wp.array(dtype=int) taxel_sensorid: wp.array[int]
qM_tiles: tuple[TileSet, ...] qM_tiles: tuple[TileSet, ...]
qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...] qLD_updates: tuple[wp.array[wp.vec3i], ...]
qLD_all_updates: wp.array(dtype=wp.vec3i) qLD_all_updates: wp.array[wp.vec3i]
qLD_level_offsets: wp.array(dtype=int) qLD_level_offsets: wp.array[int]
qM_fullm_i: wp.array(dtype=int) qM_fullm_i: wp.array[int]
qM_fullm_j: wp.array(dtype=int) qM_fullm_j: wp.array[int]
# Gather-based sparse mul_m indices (thread per DOF, no atomics) # Gather-based sparse mul_m indices (thread per DOF, no atomics)
qM_mulm_rowadr: wp.array(dtype=int) # start address for each row [nv+1] qM_mulm_rowadr: wp.array[int] # start address for each row [nv+1]
qM_mulm_col: wp.array(dtype=int) # column index to gather from qM_mulm_col: wp.array[int] # column index to gather from
qM_mulm_madr: wp.array(dtype=int) # matrix address to read qM_mulm_madr: wp.array[int] # matrix address to read
class ContactType(enum.IntFlag): class ContactType(enum.IntFlag):
@@ -1674,10 +1684,10 @@ class Constraint:
type: array("nworld", "njmax", int) type: array("nworld", "njmax", int)
id: array("nworld", "njmax", int) id: array("nworld", "njmax", int)
J_rownnz: wp.array2d(dtype=int) J_rownnz: wp.array2d[int]
J_rowadr: wp.array2d(dtype=int) J_rowadr: wp.array2d[int]
J_colind: wp.array3d(dtype=int) J_colind: wp.array3d[int]
J: wp.array3d(dtype=float) J: wp.array3d[float]
pos: array("nworld", "njmax", float) pos: array("nworld", "njmax", float)
margin: array("nworld", "njmax", float) margin: array("nworld", "njmax", float)
D: array("nworld", "njmax_pad", float) D: array("nworld", "njmax_pad", float)
@@ -1846,8 +1856,8 @@ class Data:
moment_colind: array("nworld", "nJmom", int) moment_colind: array("nworld", "nJmom", int)
actuator_moment: array("nworld", "nJmom", float) actuator_moment: array("nworld", "nJmom", float)
crb: array("nworld", "nbody", vec10) crb: array("nworld", "nbody", vec10)
qM: wp.array3d(dtype=float) qM: wp.array3d[float]
qLD: wp.array3d(dtype=float) qLD: wp.array3d[float]
qLDiagInv: array("nworld", "nv", float) qLDiagInv: array("nworld", "nv", float)
flexedge_velocity: array("nworld", "nflexedge", float) flexedge_velocity: array("nworld", "nflexedge", float)
ten_velocity: array("nworld", "ntendon", float) ten_velocity: array("nworld", "ntendon", float)
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+1
View File
@@ -57,6 +57,7 @@ dev = [
"lsprotocol>=2023.0.1,<2024.0.0", "lsprotocol>=2023.0.1,<2024.0.0",
"mujoco>=3.6.0.dev0", "mujoco>=3.6.0.dev0",
"warp-lang>=1.11.0.dev0", "warp-lang>=1.11.0.dev0",
"mjviser>=0.0.10",
] ]
# TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp
cpu = [ cpu = [
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+73 -39
View File
@@ -1,4 +1,4 @@
# Copyright 2025 The Newton Developers # Copyright 2026 The Newton Developers
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
@@ -62,6 +62,7 @@ _OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation
_KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.") _KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
_DEVICE = flags.DEFINE_string("device", None, "override the default Warp device") _DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
_REPLAY = flags.DEFINE_string("replay", None, "keyframe sequence to replay, keyframe name must prefix match") _REPLAY = flags.DEFINE_string("replay", None, "keyframe sequence to replay, keyframe name must prefix match")
_VIEWER = flags.DEFINE_enum("viewer", "mujoco", ["mujoco", "viser"], "Viewer backend (mujoco native or mjviser web)")
_VIEWER_GLOBAL_STATE = {"running": True, "step_once": False} _VIEWER_GLOBAL_STATE = {"running": True, "step_once": False}
@@ -105,6 +106,70 @@ def _compile_step(m, d):
return capture.graph return capture.graph
def _make_warp_step_fn(mjm, m, d, graph, ctrls=None):
ctrlid = 0
opt = copy.copy(mjm.opt)
def step_fn(mjm, mjd):
nonlocal ctrlid, opt, m, graph
if ctrls is not None and ctrlid < len(ctrls):
mjd.ctrl[:] = ctrls[ctrlid]
ctrlid += 1
if mjm.opt != opt:
opt = copy.copy(mjm.opt)
m = mjw.put_model(mjm)
graph = _compile_step(m, d) if wp.get_device().is_cuda else None
wp.copy(d.ctrl, wp.array([mjd.ctrl.astype(np.float32)]))
wp.copy(d.act, wp.array([mjd.act.astype(np.float32)]))
wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied.astype(np.float32)]))
wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
if graph is None:
mjw.step(m, d)
else:
wp.capture_launch(graph)
wp.synchronize()
mjw.get_data_into(mjd, mjm, d)
return step_fn
def _make_c_step_fn(ctrls=None):
if ctrls is None:
return mujoco.mj_step
ctrlid = 0
def step_fn(mjm, mjd):
nonlocal ctrlid
if ctrlid < len(ctrls):
mjd.ctrl[:] = ctrls[ctrlid]
ctrlid += 1
mujoco.mj_step(mjm, mjd)
return step_fn
def _run_viser_viewer(mjm, mjd, step_fn):
from mjviser import Viewer as MjViserViewer
MjViserViewer(mjm, mjd, step_fn=step_fn).run()
def _run_passive_viewer(mjm, mjd, step_fn):
with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
while True:
start = time.time()
if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
_VIEWER_GLOBAL_STATE["step_once"] = False
step_fn(mjm, mjd)
viewer.sync()
elapsed = time.time() - start
if elapsed < mjm.opt.timestep:
time.sleep(mjm.opt.timestep - elapsed)
def _main(argv: Sequence[str]) -> None: def _main(argv: Sequence[str]) -> None:
"""Runs viewer app.""" """Runs viewer app."""
if len(argv) < 2: if len(argv) < 2:
@@ -115,7 +180,6 @@ def _main(argv: Sequence[str]) -> None:
mjm = _load_model(epath.Path(argv[1])) mjm = _load_model(epath.Path(argv[1]))
mjd = mujoco.MjData(mjm) mjd = mujoco.MjData(mjm)
ctrls = None ctrls = None
ctrlid = 0
if _REPLAY.value: if _REPLAY.value:
keys = find_keys(mjm, _REPLAY.value) keys = find_keys(mjm, _REPLAY.value)
if not keys: if not keys:
@@ -169,45 +233,15 @@ def _main(argv: Sequence[str]) -> None:
print(f"Data\n nworld: {d.nworld} nconmax: {int(d.naconmax / d.nworld)} njmax: {d.njmax}\n") print(f"Data\n nworld: {d.nworld} nconmax: {int(d.naconmax / d.nworld)} njmax: {d.njmax}\n")
print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...") print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer: if _ENGINE.value == EngineOptions.WARP:
opt = copy.copy(mjm.opt) step_fn = _make_warp_step_fn(mjm, m, d, graph, ctrls)
while True:
start = time.time()
if ctrls is not None and ctrlid < len(ctrls):
mjd.ctrl[:] = ctrls[ctrlid]
ctrlid += 1
if _ENGINE.value == EngineOptions.C:
mujoco.mj_step(mjm, mjd)
else: # mjwarp
wp.copy(d.ctrl, wp.array([mjd.ctrl.astype(np.float32)]))
wp.copy(d.act, wp.array([mjd.act.astype(np.float32)]))
wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied.astype(np.float32)]))
wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
# if the user changed an option in the MuJoCo Simulate UI, go ahead and recompile the step
# TODO: update memory tied to option max iterations
if mjm.opt != opt:
opt = copy.copy(mjm.opt)
m = mjw.put_model(mjm)
graph = _compile_step(m, d) if wp.get_device().is_cuda else None
if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
_VIEWER_GLOBAL_STATE["step_once"] = False
if graph is None:
mjw.step(m, d)
else: else:
wp.capture_launch(graph) step_fn = _make_c_step_fn(ctrls)
wp.synchronize()
mjw.get_data_into(mjd, mjm, d)
viewer.sync() if _VIEWER.value == "viser":
_run_viser_viewer(mjm, mjd, step_fn)
elapsed = time.time() - start else:
if elapsed < mjm.opt.timestep: _run_passive_viewer(mjm, mjd, step_fn)
time.sleep(mjm.opt.timestep - elapsed)
def main(): def main():
mjx/mujoco/mjx/warp/bvh.py
+20 -19
View File
@@ -48,34 +48,35 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _refit_bvh_shim( def _refit_bvh_shim(
# Model # Model
nworld: int, nworld: int,
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
flex_elemadr: wp.array(dtype=int), flex_elemadr: wp.array[int],
flex_elemdataadr: wp.array(dtype=int), flex_elemdataadr: wp.array[int],
flex_elemnum: wp.array(dtype=int), flex_elemnum: wp.array[int],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
flex_shelldataadr: wp.array(dtype=int), flex_shelldataadr: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_vertnum: wp.array(dtype=int), flex_vertnum: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
nflex: int, nflex: int,
nflexelem: int, nflexelem: int,
# Data # Data
flexvert_xpos: wp.array2d(dtype=wp.vec3), flexvert_xpos: wp.array2d[wp.vec3],
geom_xmat: wp.array2d(dtype=wp.mat33), geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d(dtype=wp.vec3), geom_xpos: wp.array2d[wp.vec3],
# Registry # Registry
rc_id: int, rc_id: int,
# Dummy output # Dummy output
dummy: wp.array(dtype=int), dummy: wp.array[int],
): ):
_m.stat = _s _m.stat = _s
_m.opt = _o _m.opt = _o
@@ -135,7 +136,7 @@ def _refit_bvh_jax_impl(
m._impl.flex_shelldataadr, m._impl.flex_shelldataadr,
m.flex_vertadr, m.flex_vertadr,
m.flex_vertnum, m.flex_vertnum,
m.geom_dataid, jax.numpy.expand_dims(m.geom_dataid, 0),
m.geom_size, m.geom_size,
m.geom_type, m.geom_type,
m.nflex, m.nflex,
mjx/mujoco/mjx/warp/collision_driver.py
+93 -92
View File
@@ -46,67 +46,68 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _collision_shim( def _collision_shim(
# Model # Model
nworld: int, nworld: int,
block_dim: mjwp_types.BlockDim, block_dim: mjwp_types.BlockDim,
flex_conaffinity: wp.array(dtype=int), flex_conaffinity: wp.array[int],
flex_condim: wp.array(dtype=int), flex_condim: wp.array[int],
flex_contype: wp.array(dtype=int), flex_contype: wp.array[int],
flex_dim: wp.array(dtype=int), flex_dim: wp.array[int],
flex_elem: wp.array(dtype=int), flex_elem: wp.array[int],
flex_elemadr: wp.array(dtype=int), flex_elemadr: wp.array[int],
flex_elemdataadr: wp.array(dtype=int), flex_elemdataadr: wp.array[int],
flex_elemnum: wp.array(dtype=int), flex_elemnum: wp.array[int],
flex_friction: wp.array(dtype=wp.vec3), flex_friction: wp.array[wp.vec3],
flex_margin: wp.array(dtype=float), flex_margin: wp.array[float],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
flex_shell: wp.array(dtype=int), flex_shell: wp.array[int],
flex_shelldataadr: wp.array(dtype=int), flex_shelldataadr: wp.array[int],
flex_shellnum: wp.array(dtype=int), flex_shellnum: wp.array[int],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
flex_vertflexid: wp.array(dtype=int), flex_vertflexid: wp.array[int],
geom_aabb: wp.array3d(dtype=wp.vec3), geom_aabb: wp.array3d[wp.vec3],
geom_conaffinity: wp.array(dtype=int), geom_conaffinity: wp.array[int],
geom_condim: wp.array(dtype=int), geom_condim: wp.array[int],
geom_contype: wp.array(dtype=int), geom_contype: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d[wp.vec3],
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d[float],
geom_margin: wp.array2d(dtype=float), geom_margin: wp.array2d[float],
geom_pair_type_count: tuple[int, ...], geom_pair_type_count: tuple[int, ...],
geom_plugin_index: wp.array(dtype=int), geom_plugin_index: wp.array[int],
geom_priority: wp.array(dtype=int), geom_priority: wp.array[int],
geom_rbound: wp.array2d(dtype=float), geom_rbound: wp.array2d[float],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_solimp: wp.array2d(dtype=mjwp_types.vec5), geom_solimp: wp.array2d[mjwp_types.vec5],
geom_solmix: wp.array2d(dtype=float), geom_solmix: wp.array2d[float],
geom_solref: wp.array2d(dtype=wp.vec2), geom_solref: wp.array2d[wp.vec2],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
has_sdf_geom: bool, has_sdf_geom: bool,
hfield_adr: wp.array(dtype=int), hfield_adr: wp.array[int],
hfield_data: wp.array(dtype=float), hfield_data: wp.array[float],
hfield_ncol: wp.array(dtype=int), hfield_ncol: wp.array[int],
hfield_nrow: wp.array(dtype=int), hfield_nrow: wp.array[int],
hfield_size: wp.array(dtype=wp.vec4), hfield_size: wp.array[wp.vec4],
mesh_face: wp.array(dtype=wp.vec3i), mesh_face: wp.array[wp.vec3i],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
mesh_graph: wp.array(dtype=int), mesh_graph: wp.array[int],
mesh_graphadr: wp.array(dtype=int), mesh_graphadr: wp.array[int],
mesh_octadr: wp.array(dtype=int), mesh_octadr: wp.array[int],
mesh_polyadr: wp.array(dtype=int), mesh_polyadr: wp.array[int],
mesh_polymap: wp.array(dtype=int), mesh_polymap: wp.array[int],
mesh_polymapadr: wp.array(dtype=int), mesh_polymapadr: wp.array[int],
mesh_polymapnum: wp.array(dtype=int), mesh_polymapnum: wp.array[int],
mesh_polynormal: wp.array(dtype=wp.vec3), mesh_polynormal: wp.array[wp.vec3],
mesh_polynum: wp.array(dtype=int), mesh_polynum: wp.array[int],
mesh_polyvert: wp.array(dtype=int), mesh_polyvert: wp.array[int],
mesh_polyvertadr: wp.array(dtype=int), mesh_polyvertadr: wp.array[int],
mesh_polyvertnum: wp.array(dtype=int), mesh_polyvertnum: wp.array[int],
mesh_vert: wp.array(dtype=wp.vec3), mesh_vert: wp.array[wp.vec3],
mesh_vertadr: wp.array(dtype=int), mesh_vertadr: wp.array[int],
mesh_vertnum: wp.array(dtype=int), mesh_vertnum: wp.array[int],
nflex: int, nflex: int,
nflexelem: int, nflexelem: int,
nflexshelldata: int, nflexshelldata: int,
@@ -115,25 +116,25 @@ def _collision_shim(
nmaxmeshdeg: int, nmaxmeshdeg: int,
nmaxpolygon: int, nmaxpolygon: int,
nmeshface: int, nmeshface: int,
nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i), nxn_geom_pair_filtered: wp.array[wp.vec2i],
nxn_pairid: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array[wp.vec2i],
nxn_pairid_filtered: wp.array(dtype=wp.vec2i), nxn_pairid_filtered: wp.array[wp.vec2i],
oct_aabb: wp.array2d(dtype=wp.vec3), oct_aabb: wp.array2d[wp.vec3],
oct_child: wp.array(dtype=mjwp_types.vec8i), oct_child: wp.array[mjwp_types.vec8i],
oct_coeff: wp.array(dtype=mjwp_types.vec8), oct_coeff: wp.array[mjwp_types.vec8],
pair_dim: wp.array(dtype=int), pair_dim: wp.array[int],
pair_friction: wp.array2d(dtype=mjwp_types.vec5), pair_friction: wp.array2d[mjwp_types.vec5],
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d[float],
pair_margin: wp.array2d(dtype=float), pair_margin: wp.array2d[float],
pair_solimp: wp.array2d(dtype=mjwp_types.vec5), pair_solimp: wp.array2d[mjwp_types.vec5],
pair_solref: wp.array2d(dtype=wp.vec2), pair_solref: wp.array2d[wp.vec2],
pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d[wp.vec2],
plugin: wp.array(dtype=int), plugin: wp.array[int],
plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr), plugin_attr: wp.array[mjwp_types.vec_pluginattr],
opt__broadphase: int, opt__broadphase: int,
opt__broadphase_filter: int, opt__broadphase_filter: int,
opt__ccd_iterations: int, opt__ccd_iterations: int,
opt__ccd_tolerance: wp.array(dtype=float), opt__ccd_tolerance: wp.array[float],
opt__disableflags: int, opt__disableflags: int,
opt__enableflags: int, opt__enableflags: int,
opt__sdf_initpoints: int, opt__sdf_initpoints: int,
@@ -141,27 +142,27 @@ def _collision_shim(
# Data # Data
naccdmax: int, naccdmax: int,
naconmax: int, naconmax: int,
flexvert_xpos: wp.array2d(dtype=wp.vec3), flexvert_xpos: wp.array2d[wp.vec3],
geom_xmat: wp.array2d(dtype=wp.mat33), geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d(dtype=wp.vec3), geom_xpos: wp.array2d[wp.vec3],
nacon: wp.array(dtype=int), nacon: wp.array[int],
ncollision: wp.array(dtype=int), ncollision: wp.array[int],
contact__dim: wp.array(dtype=int), contact__dim: wp.array[int],
contact__dist: wp.array(dtype=float), contact__dist: wp.array[float],
contact__efc_address: wp.array2d(dtype=int), contact__efc_address: wp.array2d[int],
contact__flex: wp.array(dtype=wp.vec2i), contact__flex: wp.array[wp.vec2i],
contact__frame: wp.array(dtype=wp.mat33), contact__frame: wp.array[wp.mat33],
contact__friction: wp.array(dtype=mjwp_types.vec5), contact__friction: wp.array[mjwp_types.vec5],
contact__geom: wp.array(dtype=wp.vec2i), contact__geom: wp.array[wp.vec2i],
contact__geomcollisionid: wp.array(dtype=int), contact__geomcollisionid: wp.array[int],
contact__includemargin: wp.array(dtype=float), contact__includemargin: wp.array[float],
contact__pos: wp.array(dtype=wp.vec3), contact__pos: wp.array[wp.vec3],
contact__solimp: wp.array(dtype=mjwp_types.vec5), contact__solimp: wp.array[mjwp_types.vec5],
contact__solref: wp.array(dtype=wp.vec2), contact__solref: wp.array[wp.vec2],
contact__solreffriction: wp.array(dtype=wp.vec2), contact__solreffriction: wp.array[wp.vec2],
contact__type: wp.array(dtype=int), contact__type: wp.array[int],
contact__vert: wp.array(dtype=wp.vec2i), contact__vert: wp.array[wp.vec2i],
contact__worldid: wp.array(dtype=int), contact__worldid: wp.array[int],
): ):
_m.stat = _s _m.stat = _s
_m.opt = _o _m.opt = _o
@@ -376,7 +377,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.geom_conaffinity, m.geom_conaffinity,
m.geom_condim, m.geom_condim,
m.geom_contype, m.geom_contype,
m.geom_dataid, jax.numpy.expand_dims(m.geom_dataid, 0),
m.geom_friction, m.geom_friction,
m.geom_gap, m.geom_gap,
m.geom_margin, m.geom_margin,
mjx/mujoco/mjx/warp/forward.py
+809 -808
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/render.py
+30 -29
View File
@@ -48,42 +48,43 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _render_shim( def _render_shim(
# Model # Model
nworld: int, nworld: int,
cam_fovy: wp.array2d(dtype=float), cam_fovy: wp.array2d[float],
cam_intrinsic: wp.array2d(dtype=wp.vec4), cam_intrinsic: wp.array2d[wp.vec4],
cam_projection: wp.array(dtype=int), cam_projection: wp.array[int],
cam_sensorsize: wp.array(dtype=wp.vec2), cam_sensorsize: wp.array[wp.vec2],
flex_edge: wp.array(dtype=wp.vec2i), flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array(dtype=float), flex_radius: wp.array[float],
flex_vertadr: wp.array(dtype=int), flex_vertadr: wp.array[int],
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array2d[int],
geom_matid: wp.array2d(dtype=int), geom_matid: wp.array2d[int],
geom_rgba: wp.array2d(dtype=wp.vec4), geom_rgba: wp.array2d[wp.vec4],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
geom_type: wp.array(dtype=int), geom_type: wp.array[int],
light_active: wp.array2d(dtype=bool), light_active: wp.array2d[bool],
light_castshadow: wp.array2d(dtype=bool), light_castshadow: wp.array2d[bool],
light_type: wp.array2d(dtype=int), light_type: wp.array2d[int],
mat_rgba: wp.array2d(dtype=wp.vec4), mat_rgba: wp.array2d[wp.vec4],
mat_texid: wp.array3d(dtype=int), mat_texid: wp.array3d[int],
mat_texrepeat: wp.array2d(dtype=wp.vec2), mat_texrepeat: wp.array2d[wp.vec2],
mesh_faceadr: wp.array(dtype=int), mesh_faceadr: wp.array[int],
nlight: int, nlight: int,
# Data # Data
cam_xmat: wp.array2d(dtype=wp.mat33), cam_xmat: wp.array2d[wp.mat33],
cam_xpos: wp.array2d(dtype=wp.vec3), cam_xpos: wp.array2d[wp.vec3],
flexvert_xpos: wp.array2d(dtype=wp.vec3), flexvert_xpos: wp.array2d[wp.vec3],
geom_xmat: wp.array2d(dtype=wp.mat33), geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d(dtype=wp.vec3), geom_xpos: wp.array2d[wp.vec3],
light_xdir: wp.array2d(dtype=wp.vec3), light_xdir: wp.array2d[wp.vec3],
light_xpos: wp.array2d(dtype=wp.vec3), light_xpos: wp.array2d[wp.vec3],
# Registry # Registry
rc_id: int, rc_id: int,
rgb: wp.array2d(dtype=wp.uint32), rgb: wp.array2d[wp.uint32],
depth: wp.array2d(dtype=wp.float32), depth: wp.array2d[wp.float32],
): ):
_m.stat = _s _m.stat = _s
_m.opt = _o _m.opt = _o
@@ -164,7 +165,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
m._impl.flex_edge, m._impl.flex_edge,
m._impl.flex_radius, m._impl.flex_radius,
m.flex_vertadr, m.flex_vertadr,
m.geom_dataid, jax.numpy.expand_dims(m.geom_dataid, 0),
m.geom_matid, m.geom_matid,
m.geom_rgba, m.geom_rgba,
m.geom_size, m.geom_size,
mjx/mujoco/mjx/warp/smooth.py
+91 -90
View File
@@ -46,52 +46,53 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _kinematics_shim( def _kinematics_shim(
# Model # Model
nworld: int, nworld: int,
body_branch_start: wp.array(dtype=int), body_branch_start: wp.array[int],
body_branches: wp.array(dtype=int), body_branches: wp.array[int],
body_ipos: wp.array2d(dtype=wp.vec3), body_ipos: wp.array2d[wp.vec3],
body_iquat: wp.array2d(dtype=wp.quat), body_iquat: wp.array2d[wp.quat],
body_jntadr: wp.array(dtype=int), body_jntadr: wp.array[int],
body_jntnum: wp.array(dtype=int), body_jntnum: wp.array[int],
body_mocapid: wp.array(dtype=int), body_mocapid: wp.array[int],
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_pos: wp.array2d(dtype=wp.vec3), body_pos: wp.array2d[wp.vec3],
body_quat: wp.array2d(dtype=wp.quat), body_quat: wp.array2d[wp.quat],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_weldid: wp.array(dtype=int), body_weldid: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_pos: wp.array2d(dtype=wp.vec3), geom_pos: wp.array2d[wp.vec3],
geom_quat: wp.array2d(dtype=wp.quat), geom_quat: wp.array2d[wp.quat],
jnt_axis: wp.array2d(dtype=wp.vec3), jnt_axis: wp.array2d[wp.vec3],
jnt_pos: wp.array2d(dtype=wp.vec3), jnt_pos: wp.array2d[wp.vec3],
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
jnt_type: wp.array(dtype=int), jnt_type: wp.array[int],
nbody: int, nbody: int,
nbranch: int, nbranch: int,
ngeom: int, ngeom: int,
nsite: int, nsite: int,
qpos0: wp.array2d(dtype=float), qpos0: wp.array2d[float],
site_bodyid: wp.array(dtype=int), site_bodyid: wp.array[int],
site_pos: wp.array2d(dtype=wp.vec3), site_pos: wp.array2d[wp.vec3],
site_quat: wp.array2d(dtype=wp.quat), site_quat: wp.array2d[wp.quat],
# Data # Data
geom_xmat: wp.array2d(dtype=wp.mat33), geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d(dtype=wp.vec3), geom_xpos: wp.array2d[wp.vec3],
mocap_pos: wp.array2d(dtype=wp.vec3), mocap_pos: wp.array2d[wp.vec3],
mocap_quat: wp.array2d(dtype=wp.quat), mocap_quat: wp.array2d[wp.quat],
qpos: wp.array2d(dtype=float), qpos: wp.array2d[float],
site_xmat: wp.array2d(dtype=wp.mat33), site_xmat: wp.array2d[wp.mat33],
site_xpos: wp.array2d(dtype=wp.vec3), site_xpos: wp.array2d[wp.vec3],
xanchor: wp.array2d(dtype=wp.vec3), xanchor: wp.array2d[wp.vec3],
xaxis: wp.array2d(dtype=wp.vec3), xaxis: wp.array2d[wp.vec3],
ximat: wp.array2d(dtype=wp.mat33), ximat: wp.array2d[wp.mat33],
xipos: wp.array2d(dtype=wp.vec3), xipos: wp.array2d[wp.vec3],
xmat: wp.array2d(dtype=wp.mat33), xmat: wp.array2d[wp.mat33],
xpos: wp.array2d(dtype=wp.vec3), xpos: wp.array2d[wp.vec3],
xquat: wp.array2d(dtype=wp.quat), xquat: wp.array2d[wp.quat],
): ):
_m.stat = _s _m.stat = _s
_m.opt = _o _m.opt = _o
@@ -297,45 +298,45 @@ def kinematics_vmap(
def _tendon_shim( def _tendon_shim(
# Model # Model
nworld: int, nworld: int,
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array[int],
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array[int],
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array[int],
geom_size: wp.array2d(dtype=wp.vec3), geom_size: wp.array2d[wp.vec3],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
jnt_qposadr: wp.array(dtype=int), jnt_qposadr: wp.array[int],
ntendon: int, ntendon: int,
nwrap: int, nwrap: int,
site_bodyid: wp.array(dtype=int), site_bodyid: wp.array[int],
ten_J_colind: wp.array(dtype=int), ten_J_colind: wp.array[int],
ten_J_rowadr: wp.array(dtype=int), ten_J_rowadr: wp.array[int],
ten_J_rownnz: wp.array(dtype=int), ten_J_rownnz: wp.array[int],
tendon_adr: wp.array(dtype=int), tendon_adr: wp.array[int],
tendon_geom_adr: wp.array(dtype=int), tendon_geom_adr: wp.array[int],
tendon_jnt_adr: wp.array(dtype=int), tendon_jnt_adr: wp.array[int],
tendon_num: wp.array(dtype=int), tendon_num: wp.array[int],
tendon_site_pair_adr: wp.array(dtype=int), tendon_site_pair_adr: wp.array[int],
wrap_geom_adr: wp.array(dtype=int), wrap_geom_adr: wp.array[int],
wrap_jnt_adr: wp.array(dtype=int), wrap_jnt_adr: wp.array[int],
wrap_objid: wp.array(dtype=int), wrap_objid: wp.array[int],
wrap_prm: wp.array(dtype=float), wrap_prm: wp.array[float],
wrap_pulley_scale: wp.array(dtype=float), wrap_pulley_scale: wp.array[float],
wrap_site_pair_adr: wp.array(dtype=int), wrap_site_pair_adr: wp.array[int],
wrap_type: wp.array(dtype=int), wrap_type: wp.array[int],
# Data # Data
cdof: wp.array2d(dtype=wp.spatial_vector), cdof: wp.array2d[wp.spatial_vector],
geom_xmat: wp.array2d(dtype=wp.mat33), geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d(dtype=wp.vec3), geom_xpos: wp.array2d[wp.vec3],
qpos: wp.array2d(dtype=float), qpos: wp.array2d[float],
site_xpos: wp.array2d(dtype=wp.vec3), site_xpos: wp.array2d[wp.vec3],
subtree_com: wp.array2d(dtype=wp.vec3), subtree_com: wp.array2d[wp.vec3],
ten_J: wp.array2d(dtype=float), ten_J: wp.array2d[float],
ten_length: wp.array2d(dtype=float), ten_length: wp.array2d[float],
ten_wrapadr: wp.array2d(dtype=int), ten_wrapadr: wp.array2d[int],
ten_wrapnum: wp.array2d(dtype=int), ten_wrapnum: wp.array2d[int],
wrap_obj: wp.array2d(dtype=wp.vec2i), wrap_obj: wp.array2d[wp.vec2i],
wrap_xpos: wp.array2d(dtype=wp.spatial_vector), wrap_xpos: wp.array2d[wp.spatial_vector],
): ):
_m.stat = _s _m.stat = _s
_m.opt = _o _m.opt = _o
@@ -489,26 +490,26 @@ def tendon_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
def _com_pos_shim( def _com_pos_shim(
# Model # Model
nworld: int, nworld: int,
body_inertia: wp.array2d(dtype=wp.vec3), body_inertia: wp.array2d[wp.vec3],
body_mass: wp.array2d(dtype=float), body_mass: wp.array2d[float],
body_parentid: wp.array(dtype=int), body_parentid: wp.array[int],
body_rootid: wp.array(dtype=int), body_rootid: wp.array[int],
body_subtreemass: wp.array2d(dtype=float), body_subtreemass: wp.array2d[float],
body_tree: tuple[wp.array(dtype=int), ...], body_tree: tuple[wp.array[int], ...],
jnt_bodyid: wp.array(dtype=int), jnt_bodyid: wp.array[int],
jnt_dofadr: wp.array(dtype=int), jnt_dofadr: wp.array[int],
jnt_type: wp.array(dtype=int), jnt_type: wp.array[int],
nbody: int, nbody: int,
njnt: int, njnt: int,
# Data # Data
cdof: wp.array2d(dtype=wp.spatial_vector), cdof: wp.array2d[wp.spatial_vector],
cinert: wp.array2d(dtype=mjwp_types.vec10), cinert: wp.array2d[mjwp_types.vec10],
subtree_com: wp.array2d(dtype=wp.vec3), subtree_com: wp.array2d[wp.vec3],
xanchor: wp.array2d(dtype=wp.vec3), xanchor: wp.array2d[wp.vec3],
xaxis: wp.array2d(dtype=wp.vec3), xaxis: wp.array2d[wp.vec3],
ximat: wp.array2d(dtype=wp.mat33), ximat: wp.array2d[wp.mat33],
xipos: wp.array2d(dtype=wp.vec3), xipos: wp.array2d[wp.vec3],
xmat: wp.array2d(dtype=wp.mat33), xmat: wp.array2d[wp.mat33],
): ):
_m.stat = _s _m.stat = _s
_m.opt = _o _m.opt = _o
mjx/mujoco/mjx/warp/types.py
+2 -2
View File
@@ -672,7 +672,7 @@ _NDIM = {
'geom_conaffinity': 1, 'geom_conaffinity': 1,
'geom_condim': 1, 'geom_condim': 1,
'geom_contype': 1, 'geom_contype': 1,
'geom_dataid': 1, 'geom_dataid': 2,
'geom_fluid': 2, 'geom_fluid': 2,
'geom_friction': 3, 'geom_friction': 3,
'geom_gap': 2, 'geom_gap': 2,
@@ -1238,7 +1238,7 @@ _BATCH_DIM = {
'geom_conaffinity': False, 'geom_conaffinity': False,
'geom_condim': False, 'geom_condim': False,
'geom_contype': False, 'geom_contype': False,
'geom_dataid': False, 'geom_dataid': True,
'geom_fluid': False, 'geom_fluid': False,
'geom_friction': True, 'geom_friction': True,
'geom_gap': True, 'geom_gap': True,