From 8bf11166b1403a36d1ca7c66298822f8d9950423 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Tue, 14 Apr 2026 12:55:06 -0700 Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub. PiperOrigin-RevId: 899731493 Change-Id: Ib563023047e85db1aff56e78724dc28d769ef91e --- mjx/mujoco/mjx/_src/io.py | 3 + .../third_party/mujoco_warp/_src/benchmark.py | 12 +- .../mujoco_warp/_src/block_cholesky.py | 10 +- .../mjx/third_party/mujoco_warp/_src/bvh.py | 148 +- .../mujoco_warp/_src/collision_convex.py | 377 ++-- .../mujoco_warp/_src/collision_core.py | 126 +- .../mujoco_warp/_src/collision_driver.py | 108 +- .../mujoco_warp/_src/collision_flex.py | 276 +-- .../mujoco_warp/_src/collision_gjk.py | 142 +- .../mujoco_warp/_src/collision_primitive.py | 496 ++--- .../mujoco_warp/_src/collision_sdf.py | 182 +- .../mujoco_warp/_src/constraint.py | 800 ++++---- .../mujoco_warp/_src/derivative.py | 106 +- .../third_party/mujoco_warp/_src/forward.py | 226 +-- .../third_party/mujoco_warp/_src/inverse.py | 20 +- .../mjx/third_party/mujoco_warp/_src/io.py | 312 ++-- .../third_party/mujoco_warp/_src/island.py | 44 +- .../third_party/mujoco_warp/_src/passive.py | 174 +- .../mjx/third_party/mujoco_warp/_src/ray.py | 226 +-- .../third_party/mujoco_warp/_src/render.py | 216 +-- .../mujoco_warp/_src/render_util.py | 28 +- .../third_party/mujoco_warp/_src/sensor.py | 884 ++++----- .../third_party/mujoco_warp/_src/smooth.py | 942 +++++----- .../third_party/mujoco_warp/_src/solver.py | 648 +++---- .../third_party/mujoco_warp/_src/support.py | 196 +- .../mjx/third_party/mujoco_warp/_src/types.py | 122 +- .../third_party/mujoco_warp/pyproject.toml | 1 + .../mjx/third_party/mujoco_warp/viewer.py | 114 +- mjx/mujoco/mjx/warp/bvh.py | 39 +- mjx/mujoco/mjx/warp/collision_driver.py | 185 +- mjx/mujoco/mjx/warp/forward.py | 1617 +++++++++-------- mjx/mujoco/mjx/warp/render.py | 59 +- mjx/mujoco/mjx/warp/smooth.py | 181 +- mjx/mujoco/mjx/warp/types.py | 4 +- 34 files changed, 4556 insertions(+), 4468 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index d9ef50b3..74937bc0 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -456,6 +456,9 @@ def _put_model_warp( if not hasattr(mw, k) or k in ('stat', 'opt'): continue 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 impl_fields = {} diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py index 5712b81f..293aca63 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py @@ -41,18 +41,18 @@ def _sum(stack1, stack2): @wp.kernel def ctrl_noise( # Model: - opt_timestep: wp.array(dtype=float), - actuator_ctrllimited: wp.array(dtype=bool), - actuator_ctrlrange: wp.array2d(dtype=wp.vec2), + opt_timestep: wp.array[float], + actuator_ctrllimited: wp.array[bool], + actuator_ctrlrange: wp.array2d[wp.vec2], # Data in: - ctrl_in: wp.array2d(dtype=float), + ctrl_in: wp.array2d[float], # In: - ctrl_center: wp.array1d(dtype=float), + ctrl_center: wp.array[float], step: int, ctrlnoisestd: float, ctrlnoiserate: float, # Data out: - ctrl_out: wp.array2d(dtype=float), + ctrl_out: wp.array2d[float], ): worldid, actid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py index e354000a..628ad5a0 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py @@ -23,10 +23,10 @@ def create_blocked_cholesky_func(block_size: int): @wp.func def blocked_cholesky_func( # In: - A: wp.array(dtype=float, ndim=2), + A: wp.array2d[float], matrix_size: int, # 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. @@ -68,11 +68,11 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int) @wp.func def blocked_cholesky_solve_func( # In: - L: wp.array(dtype=float, ndim=2), - b: wp.array(dtype=float, ndim=2), + L: wp.array2d[float], + b: wp.array2d[float], matrix_size: int, # Out: - x: wp.array(dtype=float, ndim=2), + x: wp.array2d[float], ): """Block Cholesky factorization and solve. diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py index 58f7b221..32899e07 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py @@ -173,21 +173,21 @@ def _compute_cylinder_bounds( @wp.kernel def _compute_bvh_bounds( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: bvh_ngeom: int, - enabled_geom_ids: wp.array(dtype=int), - mesh_bounds_size: wp.array(dtype=wp.vec3), - hfield_bounds_size: wp.array(dtype=wp.vec3), + enabled_geom_ids: wp.array[int], + mesh_bounds_size: wp.array[wp.vec3], + hfield_bounds_size: wp.array[wp.vec3], # Out: - lower_out: wp.array(dtype=wp.vec3), - upper_out: wp.array(dtype=wp.vec3), - group_out: wp.array(dtype=int), + lower_out: wp.array[wp.vec3], + upper_out: wp.array[wp.vec3], + group_out: wp.array[int], ): worldid, geom_local_id = wp.tid() geom_id = enabled_geom_ids[geom_local_id] @@ -205,8 +205,13 @@ def _compute_bvh_bounds( elif type == GeomType.PLANE: lower_bound, upper_bound = _compute_plane_bounds(pos, rot, size) elif type == GeomType.MESH: - size = mesh_bounds_size[geom_dataid[geom_id]] - lower_bound, upper_bound = _compute_box_bounds(pos, rot, size) + 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) + else: + lower_bound = pos + upper_bound = pos elif type == GeomType.ELLIPSOID: lower_bound, upper_bound = _compute_ellipsoid_bounds(pos, rot, size) elif type == GeomType.CYLINDER: @@ -214,9 +219,14 @@ def _compute_bvh_bounds( elif type == GeomType.BOX: lower_bound, upper_bound = _compute_box_bounds(pos, rot, size) elif type == GeomType.HFIELD: - size = hfield_bounds_size[geom_dataid[geom_id]] - hfield_center = pos + rot[:, 2] * size[2] - lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size) + 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] + 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 upper_out[worldid * bvh_ngeom + geom_local_id] = upper_bound @@ -228,7 +238,7 @@ def compute_bvh_group_roots( # In: bvh_id: wp.uint64, # Out: - group_root_out: wp.array(dtype=int), + group_root_out: wp.array[int], ): tid = wp.tid() root = wp.bvh_get_group_root(bvh_id, tid) @@ -238,21 +248,21 @@ def compute_bvh_group_roots( @wp.kernel def _compute_flex_bvh_bounds( # Model: - flex_vertadr: wp.array(dtype=int), - flex_vertnum: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_radius: wp.array(dtype=float), + flex_vertadr: wp.array[int], + flex_vertnum: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # In: - flex_geom_flexid: wp.array(dtype=int), - flex_geom_edgeid: wp.array(dtype=int), + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], bvh_ngeom: int, total_bvh_size: int, # Out: - lower_out: wp.array(dtype=wp.vec3), - upper_out: wp.array(dtype=wp.vec3), - group_out: wp.array(dtype=int), + lower_out: wp.array[wp.vec3], + upper_out: wp.array[wp.vec3], + group_out: wp.array[int], ): 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 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_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) @@ -599,16 +609,16 @@ def build_hfield_bvh( def accumulate_flex_vertex_normals( # Model: nflex: int, - flex_dim: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elem: wp.array[int], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # 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.""" worldid, elemid = wp.tid() @@ -644,7 +654,7 @@ def accumulate_flex_vertex_normals( @wp.kernel def normalize_vertex_normals( # Out: - flexvert_norm_out: wp.array2d(dtype=wp.vec3), + flexvert_norm_out: wp.array2d[wp.vec3], ): """Normalize accumulated vertex normals.""" worldid, vertid = wp.tid() @@ -654,20 +664,20 @@ def normalize_vertex_normals( @wp.kernel def _build_flex_2d_elements( # Model: - flex_elem: wp.array(dtype=int), + flex_elem: wp.array[int], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # In: - flexvert_norm_in: wp.array2d(dtype=wp.vec3), + flexvert_norm_in: wp.array2d[wp.vec3], elem_adr: int, vert_adr: int, face_offset: int, radius: float, nfaces: int, # Out: - face_point_out: wp.array(dtype=wp.vec3), - face_index_out: wp.array(dtype=int), - group_out: wp.array(dtype=int), + face_point_out: wp.array[wp.vec3], + face_index_out: wp.array[int], + group_out: wp.array[int], ): """Create faces from 2D flex elements (triangles). @@ -728,20 +738,20 @@ def _build_flex_2d_elements( @wp.kernel def _build_flex_2d_sides( # Model: - flex_shell: wp.array(dtype=int), + flex_shell: wp.array[int], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # In: - flexvert_norm_in: wp.array2d(dtype=wp.vec3), + flexvert_norm_in: wp.array2d[wp.vec3], shell_adr: int, vert_adr: int, face_offset: int, radius: float, nface: int, # Out: - face_point_out: wp.array(dtype=wp.vec3), - face_index_out: wp.array(dtype=int), - group_out: wp.array(dtype=int), + face_point_out: wp.array[wp.vec3], + face_index_out: wp.array[int], + group_out: wp.array[int], ): """Create side faces from 2D flex shell fragments. @@ -790,18 +800,18 @@ def _build_flex_2d_sides( @wp.kernel def _build_flex_3d_shells( # Model: - flex_shell: wp.array(dtype=int), + flex_shell: wp.array[int], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # In: shell_adr: int, vert_adr: int, face_offset: int, nface: int, # Out: - face_point_out: wp.array(dtype=wp.vec3), - face_index_out: wp.array(dtype=int), - group_out: wp.array(dtype=int), + face_point_out: wp.array[wp.vec3], + face_index_out: wp.array[int], + group_out: wp.array[int], ): """Create faces from 3D flex shell fragments (triangles). @@ -836,22 +846,22 @@ def _build_flex_3d_shells( @wp.kernel def _update_flex_2d_face_points( # Model: - flex_vertadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), - flex_shell: wp.array(dtype=int), - flex_radius: wp.array(dtype=float), + flex_vertadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_elem: wp.array[int], + flex_shell: wp.array[int], + flex_radius: wp.array[float], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # In: - flexvert_norm_in: wp.array2d(dtype=wp.vec3), + flexvert_norm_in: wp.array2d[wp.vec3], flex_id: int, nface: int, smooth: bool, # Out: - face_point_out: wp.array(dtype=wp.vec3), + face_point_out: wp.array[wp.vec3], ): worldid, workid = wp.tid() @@ -935,16 +945,16 @@ def _update_flex_2d_face_points( @wp.kernel def _update_flex_3d_face_points( # Model: - flex_vertadr: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_shell: wp.array(dtype=int), + flex_vertadr: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_shell: wp.array[int], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # In: flex_id: int, nface: int, # Out: - face_point_out: wp.array(dtype=wp.vec3), + face_point_out: wp.array[wp.vec3], ): worldid, shellid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py index 12d9824a..4f6b21c5 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py @@ -55,20 +55,20 @@ mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float) @wp.func def _hfield_filter( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array2d(dtype=float), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_graph: wp.array(dtype=int), - hfield_size: wp.array(dtype=wp.vec4), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + geom_rbound: wp.array2d[float], + geom_margin: wp.array2d[float], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_graph: wp.array[int], + hfield_size: wp.array[wp.vec4], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, g1: int, @@ -79,7 +79,8 @@ def _hfield_filter( See MuJoCo mjc_ConvexHField. """ # height field info - hfdataid = geom_dataid[g1] + dataid_setid = worldid % geom_dataid.shape[0] + hfdataid = geom_dataid[dataid_setid, g1] size1 = hfield_size[hfdataid] # geom info @@ -124,7 +125,7 @@ def _hfield_filter( # load mesh vertex data for support function queries 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.vertnum = wp.where(dataid >= 0, mesh_vertnum[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) def ccd_hfield_kernel( # Model: - opt_ccd_tolerance: wp.array(dtype=float), - geom_type: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rbound: wp.array2d(dtype=float), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_graph: wp.array(dtype=int), - mesh_polynum: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - hfield_size: wp.array(dtype=wp.vec4), - hfield_nrow: wp.array(dtype=int), - hfield_ncol: wp.array(dtype=int), - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), + opt_ccd_tolerance: wp.array[float], + geom_type: wp.array[int], + geom_condim: wp.array[int], + geom_dataid: wp.array2d[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_size: wp.array2d[wp.vec3], + geom_rbound: wp.array2d[float], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_graph: wp.array[int], + mesh_polynum: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polymap: wp.array[int], + hfield_size: wp.array[wp.vec4], + hfield_nrow: wp.array[int], + hfield_ncol: wp.array[int], + hfield_adr: wp.array[int], + hfield_data: wp.array[float], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + pair_friction: wp.array2d[vec5], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], naconmax_in: int, naccdmax_in: int, - ncollision_in: wp.array(dtype=int), + ncollision_in: wp.array[int], # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), - collision_worldid_in: wp.array(dtype=int), - epa_vert_in: wp.array2d(dtype=wp.vec3), - epa_vert_index_in: wp.array2d(dtype=int), - epa_face_in: wp.array2d(dtype=int), - epa_pr_in: wp.array2d(dtype=wp.vec3), - epa_norm2_in: wp.array2d(dtype=float), - epa_horizon_in: wp.array2d(dtype=int), - nccd_in: wp.array(dtype=int), + collision_pair_in: wp.array[wp.vec2i], + collision_pairid_in: wp.array[wp.vec2i], + collision_worldid_in: wp.array[int], + epa_vert_in: wp.array2d[wp.vec3], + epa_vert_index_in: wp.array2d[int], + epa_face_in: wp.array2d[int], + epa_pr_in: wp.array2d[wp.vec3], + epa_norm2_in: wp.array2d[float], + epa_horizon_in: wp.array2d[int], + nccd_in: wp.array[int], # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): collisionid = wp.tid() if collisionid >= ncollision_in[0]: @@ -339,7 +340,7 @@ def ccd_hfield_kernel_builder( geom1.rot = wp.identity(n=3, dtype=float) # see MuJoCo mjc_ConvexHField - geom1_dataid = geom_dataid[g1] + geom1_dataid = geom_dataid[worldid % geom_dataid.shape[0], g1] # height field subgrid nrow = hfield_nrow[geom1_dataid] @@ -710,27 +711,27 @@ def ccd_kernel_builder( @wp.func def eval_ccd_write_contact( # Model: - opt_ccd_tolerance: wp.array(dtype=float), + opt_ccd_tolerance: wp.array[float], # Data in: naconmax_in: int, # In: - epa_vert_in: wp.array2d(dtype=wp.vec3), - epa_vert_index_in: wp.array2d(dtype=int), - epa_face_in: wp.array2d(dtype=int), - epa_pr_in: wp.array2d(dtype=wp.vec3), - epa_norm2_in: wp.array2d(dtype=float), - epa_horizon_in: wp.array2d(dtype=int), - multiccd_polygon_in: wp.array2d(dtype=wp.vec3), - multiccd_clipped_in: wp.array2d(dtype=wp.vec3), - multiccd_pnormal_in: wp.array2d(dtype=wp.vec3), - multiccd_pdist_in: wp.array2d(dtype=float), - multiccd_idx1_in: wp.array2d(dtype=int), - multiccd_idx2_in: wp.array2d(dtype=int), - multiccd_n1_in: wp.array2d(dtype=wp.vec3), - multiccd_n2_in: wp.array2d(dtype=wp.vec3), - multiccd_endvert_in: wp.array2d(dtype=wp.vec3), - multiccd_face1_in: wp.array2d(dtype=wp.vec3), - multiccd_face2_in: wp.array2d(dtype=wp.vec3), + epa_vert_in: wp.array2d[wp.vec3], + epa_vert_index_in: wp.array2d[int], + epa_face_in: wp.array2d[int], + epa_pr_in: wp.array2d[wp.vec3], + epa_norm2_in: wp.array2d[float], + epa_horizon_in: wp.array2d[int], + multiccd_polygon_in: wp.array2d[wp.vec3], + multiccd_clipped_in: wp.array2d[wp.vec3], + multiccd_pnormal_in: wp.array2d[wp.vec3], + multiccd_pdist_in: wp.array2d[float], + multiccd_idx1_in: wp.array2d[int], + multiccd_idx2_in: wp.array2d[int], + multiccd_n1_in: wp.array2d[wp.vec3], + multiccd_n2_in: wp.array2d[wp.vec3], + multiccd_endvert_in: wp.array2d[wp.vec3], + multiccd_face1_in: wp.array2d[wp.vec3], + multiccd_face2_in: wp.array2d[wp.vec3], geom1: Geom, geom2: Geom, geoms: wp.vec2i, @@ -747,21 +748,21 @@ def ccd_kernel_builder( x2: wp.vec3, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ) -> int: points = mat43() witness1 = mat43() @@ -892,83 +893,83 @@ def ccd_kernel_builder( @wp.kernel(module="unique", enable_backward=False) def ccd_kernel( # Model: - opt_ccd_tolerance: wp.array(dtype=float), - geom_type: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_graph: wp.array(dtype=int), - mesh_polynum: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), + opt_ccd_tolerance: wp.array[float], + geom_type: wp.array[int], + geom_condim: wp.array[int], + geom_dataid: wp.array2d[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_size: wp.array2d[wp.vec3], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_graph: wp.array[int], + mesh_polynum: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polymap: wp.array[int], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + pair_friction: wp.array2d[vec5], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], naconmax_in: int, naccdmax_in: int, - ncollision_in: wp.array(dtype=int), + ncollision_in: wp.array[int], # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), - collision_worldid_in: wp.array(dtype=int), - epa_vert_in: wp.array2d(dtype=wp.vec3), - epa_vert_index_in: wp.array2d(dtype=int), - epa_face_in: wp.array2d(dtype=int), - epa_pr_in: wp.array2d(dtype=wp.vec3), - epa_norm2_in: wp.array2d(dtype=float), - epa_horizon_in: wp.array2d(dtype=int), - multiccd_polygon_in: wp.array2d(dtype=wp.vec3), - multiccd_clipped_in: wp.array2d(dtype=wp.vec3), - multiccd_pnormal_in: wp.array2d(dtype=wp.vec3), - multiccd_pdist_in: wp.array2d(dtype=float), - multiccd_idx1_in: wp.array2d(dtype=int), - multiccd_idx2_in: wp.array2d(dtype=int), - multiccd_n1_in: wp.array2d(dtype=wp.vec3), - multiccd_n2_in: wp.array2d(dtype=wp.vec3), - multiccd_endvert_in: wp.array2d(dtype=wp.vec3), - multiccd_face1_in: wp.array2d(dtype=wp.vec3), - multiccd_face2_in: wp.array2d(dtype=wp.vec3), - nccd_in: wp.array(dtype=int), + collision_pair_in: wp.array[wp.vec2i], + collision_pairid_in: wp.array[wp.vec2i], + collision_worldid_in: wp.array[int], + epa_vert_in: wp.array2d[wp.vec3], + epa_vert_index_in: wp.array2d[int], + epa_face_in: wp.array2d[int], + epa_pr_in: wp.array2d[wp.vec3], + epa_norm2_in: wp.array2d[float], + epa_horizon_in: wp.array2d[int], + multiccd_polygon_in: wp.array2d[wp.vec3], + multiccd_clipped_in: wp.array2d[wp.vec3], + multiccd_pnormal_in: wp.array2d[wp.vec3], + multiccd_pdist_in: wp.array2d[float], + multiccd_idx1_in: wp.array2d[int], + multiccd_idx2_in: wp.array2d[int], + multiccd_n1_in: wp.array2d[wp.vec3], + multiccd_n2_in: wp.array2d[wp.vec3], + multiccd_endvert_in: wp.array2d[wp.vec3], + multiccd_face1_in: wp.array2d[wp.vec3], + multiccd_face2_in: wp.array2d[wp.vec3], + nccd_in: wp.array[int], # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): collisionid = wp.tid() if collisionid >= ncollision_in[0]: diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py index b7affa7d..96da38e0 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py @@ -47,44 +47,44 @@ class Geom: hfprism: mat63 vertadr: int vertnum: int - vert: wp.array(dtype=wp.vec3) + vert: wp.array[wp.vec3] graphadr: int - graph: wp.array(dtype=int) + graph: wp.array[int] mesh_polynum: int mesh_polyadr: int - mesh_polynormal: wp.array(dtype=wp.vec3) - mesh_polyvertadr: wp.array(dtype=int) - mesh_polyvertnum: wp.array(dtype=int) - mesh_polyvert: wp.array(dtype=int) - mesh_polymapadr: wp.array(dtype=int) - mesh_polymapnum: wp.array(dtype=int) - mesh_polymap: wp.array(dtype=int) + mesh_polynormal: wp.array[wp.vec3] + mesh_polyvertadr: wp.array[int] + mesh_polyvertnum: wp.array[int] + mesh_polyvert: wp.array[int] + mesh_polymapadr: wp.array[int] + mesh_polymapnum: wp.array[int] + mesh_polymap: wp.array[int] index: int @wp.func def geom_collision_pair( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_graph: wp.array(dtype=int), - mesh_polynum: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_graph: wp.array[int], + mesh_polynum: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polymap: wp.array[int], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: geoms: wp.vec2i, worldid: int, @@ -109,8 +109,10 @@ def geom_collision_pair( # 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]) + dataid_setid = worldid % geom_dataid.shape[0] + 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.vertnum = wp.where(dataid >= 0, mesh_vertnum[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 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.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1) geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1) @@ -174,21 +176,21 @@ def write_contact( pairid_in: wp.vec2i, worldid_in: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ) -> int: """Atomically write a detected contact into the contact output arrays. @@ -233,24 +235,24 @@ def write_contact( @wp.func def contact_params( # Model: - geom_condim: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), + geom_condim: wp.array[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + pair_friction: wp.array2d[vec5], # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), + collision_pair_in: wp.array[wp.vec2i], + collision_pairid_in: wp.array[wp.vec2i], cid: int, worldid: int, ): diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py index 786ee9b8..c87b21f1 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py @@ -80,8 +80,8 @@ MJ_COLLISION_TABLE = { @wp.kernel def _zero_nacon_ncollision( # Data out: - nacon_out: wp.array(dtype=int), - ncollision_out: wp.array(dtype=int), + nacon_out: wp.array[int], + ncollision_out: wp.array[int], ): ncollision_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 def func( # Model: - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array2d(dtype=float), + geom_aabb: wp.array3d[wp.vec3], + geom_rbound: wp.array2d[float], + geom_margin: wp.array2d[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: geom1: int, geom2: int, @@ -324,8 +324,8 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound @wp.func def _add_geom_pair( # Model: - geom_type: wp.array(dtype=int), - nxn_pairid: wp.array(dtype=wp.vec2i), + geom_type: wp.array[int], + nxn_pairid: wp.array[wp.vec2i], # Data in: naconmax_in: int, # In: @@ -334,11 +334,11 @@ def _add_geom_pair( worldid: int, nxnid: int, # Data out: - ncollision_out: wp.array(dtype=int), + ncollision_out: wp.array[int], # Out: - collision_pair_out: wp.array(dtype=wp.vec2i), - collision_pairid_out: wp.array(dtype=wp.vec2i), - collision_worldid_out: wp.array(dtype=int), + collision_pair_out: wp.array[wp.vec2i], + collision_pairid_out: wp.array[wp.vec2i], + collision_worldid_out: wp.array[int], ): pairid = wp.atomic_add(ncollision_out, 0, 1) @@ -359,7 +359,7 @@ def _add_geom_pair( @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: mid = (lower + upper) >> 1 if values[mid] > value: @@ -375,18 +375,18 @@ def _sap_project(opt_broadphase: int): def sap_project( # Model: ngeom: int, - geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array2d(dtype=float), + geom_rbound: wp.array2d[float], + geom_margin: wp.array2d[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], nworld_in: int, # In: direction_in: wp.vec3, # Out: - projection_lower_out: wp.array2d(dtype=float), - projection_upper_out: wp.array2d(dtype=float), - sort_index_out: wp.array2d(dtype=int), - segmented_index_out: wp.array(dtype=int), + projection_lower_out: wp.array2d[float], + projection_upper_out: wp.array2d[float], + sort_index_out: wp.array2d[int], + segmented_index_out: wp.array[int], ): worldid, geomid = wp.tid() @@ -422,11 +422,11 @@ def _sap_range( # Model: ngeom: int, # In: - projection_lower_in: wp.array2d(dtype=float), - projection_upper_in: wp.array2d(dtype=float), - sort_index_in: wp.array2d(dtype=int), + projection_lower_in: wp.array2d[float], + projection_upper_in: wp.array2d[float], + sort_index_in: wp.array2d[int], # Out: - range_out: wp.array2d(dtype=int), + range_out: wp.array2d[int], ): worldid, geomid = wp.tid() @@ -448,26 +448,26 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i def kernel( # Model: ngeom: int, - geom_type: wp.array(dtype=int), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array2d(dtype=float), - nxn_pairid: wp.array(dtype=wp.vec2i), + geom_type: wp.array[int], + geom_aabb: wp.array3d[wp.vec3], + geom_rbound: wp.array2d[float], + geom_margin: wp.array2d[float], + nxn_pairid: wp.array[wp.vec2i], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], nworld_in: int, naconmax_in: int, # In: - sort_index_in: wp.array2d(dtype=int), - cumulative_sum_in: wp.array(dtype=int), + sort_index_in: wp.array2d[int], + cumulative_sum_in: wp.array[int], nsweep_in: int, # Data out: - ncollision_out: wp.array(dtype=int), + ncollision_out: wp.array[int], # Out: - collision_pair_out: wp.array(dtype=wp.vec2i), - collision_pairid_out: wp.array(dtype=wp.vec2i), - collision_worldid_out: wp.array(dtype=int), + collision_pair_out: wp.array[wp.vec2i], + collision_pairid_out: wp.array[wp.vec2i], + collision_worldid_out: wp.array[int], ): worldgeomid = wp.tid() @@ -528,11 +528,11 @@ def _segmented_sort(tile_size: int): @wp.kernel(module="unique") def segmented_sort( # In: - projection_lower_in: wp.array2d(dtype=float), - sort_index_in: wp.array2d(dtype=int), + projection_lower_in: wp.array2d[float], + sort_index_in: wp.array2d[int], # Out: - projection_lower_out: wp.array2d(dtype=float), - sort_index_out: wp.array2d(dtype=int), + projection_lower_out: wp.array2d[float], + sort_index_out: wp.array2d[int], ): 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) def kernel( # Model: - geom_type: wp.array(dtype=int), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array2d(dtype=float), - nxn_geom_pair: wp.array(dtype=wp.vec2i), - nxn_pairid: wp.array(dtype=wp.vec2i), + geom_type: wp.array[int], + geom_aabb: wp.array3d[wp.vec3], + geom_rbound: wp.array2d[float], + geom_margin: wp.array2d[float], + nxn_geom_pair: wp.array[wp.vec2i], + nxn_pairid: wp.array[wp.vec2i], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], naconmax_in: int, # Data out: - ncollision_out: wp.array(dtype=int), + ncollision_out: wp.array[int], # Out: - collision_pair_out: wp.array(dtype=wp.vec2i), - collision_pairid_out: wp.array(dtype=wp.vec2i), - collision_worldid_out: wp.array(dtype=int), + collision_pair_out: wp.array[wp.vec2i], + collision_pairid_out: wp.array[wp.vec2i], + collision_worldid_out: wp.array[int], ): worldid, elementid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py index 215423e9..cd5c9087 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py @@ -46,22 +46,22 @@ def _write_flex_contact( vertid: int, worldid: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_flex_out: wp.array(dtype=wp.vec2i), - contact_vert_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_flex_out: wp.array[wp.vec2i], + contact_vert_out: wp.array[wp.vec2i], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): if dist >= margin or dist >= MJ_MAXVAL: return @@ -110,22 +110,22 @@ def _collide_geom_triangle( vertex_id: int, worldid: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_flex_out: wp.array(dtype=wp.vec2i), - contact_vert_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_flex_out: wp.array[wp.vec2i], + contact_vert_out: wp.array[wp.vec2i], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): if gtype == int(GeomType.SPHERE): sphere_radius = size_val[0] @@ -262,41 +262,41 @@ def _flex_plane_narrowphase( # Model: ngeom: int, nflexvert: int, - geom_type: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - flex_condim: wp.array(dtype=int), - flex_friction: wp.array(dtype=wp.vec3), - flex_margin: wp.array(dtype=float), - flex_vertadr: wp.array(dtype=int), - flex_radius: wp.array(dtype=float), - flex_vertflexid: wp.array(dtype=int), + geom_type: wp.array[int], + geom_condim: wp.array[int], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + flex_condim: wp.array[int], + flex_friction: wp.array[wp.vec3], + flex_margin: wp.array[float], + flex_vertadr: wp.array[int], + flex_radius: wp.array[float], + flex_vertflexid: wp.array[int], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], nworld_in: int, naconmax_in: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_flex_out: wp.array(dtype=wp.vec2i), - contact_vert_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_flex_out: wp.array[wp.vec2i], + contact_vert_out: wp.array[wp.vec2i], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): worldid, vertid = wp.tid() @@ -382,48 +382,48 @@ def _flex_narrowphase_dim2( # Model: ngeom: int, nflex: int, - geom_type: wp.array(dtype=int), - geom_contype: wp.array(dtype=int), - geom_conaffinity: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - flex_contype: wp.array(dtype=int), - flex_conaffinity: wp.array(dtype=int), - flex_margin: wp.array(dtype=float), - flex_dim: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), - flex_radius: wp.array(dtype=float), + geom_type: wp.array[int], + geom_contype: wp.array[int], + geom_conaffinity: wp.array[int], + geom_condim: wp.array[int], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_size: wp.array2d[wp.vec3], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + flex_contype: wp.array[int], + flex_conaffinity: wp.array[int], + flex_margin: wp.array[float], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elem: wp.array[int], + flex_radius: wp.array[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], nworld_in: int, naconmax_in: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_flex_out: wp.array(dtype=wp.vec2i), - contact_vert_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_flex_out: wp.array[wp.vec2i], + contact_vert_out: wp.array[wp.vec2i], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): worldid, elemid = wp.tid() @@ -533,47 +533,47 @@ def _flex_narrowphase_dim3( # Model: ngeom: int, nflex: int, - geom_type: wp.array(dtype=int), - geom_contype: wp.array(dtype=int), - geom_conaffinity: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - flex_contype: wp.array(dtype=int), - flex_conaffinity: wp.array(dtype=int), - flex_margin: wp.array(dtype=float), - flex_dim: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_shellnum: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_shell: wp.array(dtype=int), - flex_radius: wp.array(dtype=float), + geom_type: wp.array[int], + geom_contype: wp.array[int], + geom_conaffinity: wp.array[int], + geom_condim: wp.array[int], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_size: wp.array2d[wp.vec3], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + flex_contype: wp.array[int], + flex_conaffinity: wp.array[int], + flex_margin: wp.array[float], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_shellnum: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_shell: wp.array[int], + flex_radius: wp.array[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], nworld_in: int, naconmax_in: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_flex_out: wp.array(dtype=wp.vec2i), - contact_vert_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_flex_out: wp.array[wp.vec2i], + contact_vert_out: wp.array[wp.vec2i], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): worldid, shellid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py index fe1c4445..21e61b7c 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py @@ -63,20 +63,20 @@ class Polytope: status: int # vertices in polytope (packed geom1 followed by geom2) - vert: wp.array(dtype=wp.vec3) - vert_index: wp.array(dtype=int) + vert: wp.array[wp.vec3] + vert_index: wp.array[int] nvert: int # faces in polytope # 10 bits per each vertex index, while the last significant bits are for # invalid and deleted face - face: wp.array(dtype=int) - face_pr: wp.array(dtype=wp.vec3) - face_norm2: wp.array(dtype=float) + face: wp.array[int] + face_pr: wp.array[wp.vec3] + face_norm2: wp.array[float] nface: int # edges that make up the horizon when adding new vertices to polytope - horizon: wp.array(dtype=int) + horizon: wp.array[int] nhorizon: int @@ -1335,7 +1335,7 @@ def _area4(a: wp.vec3, b: wp.vec3, c: wp.vec3, d: wp.vec3) -> float: @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).""" b = int(1) 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 @wp.func 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]: v1i = vert_index[2 * face[0] + 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 @wp.func -def _aligned_faces( - vert1: wp.array(dtype=wp.vec3), len1: int, vert2: wp.array(dtype=wp.vec3), len2: int -) -> Tuple[int, wp.vec2i]: +def _aligned_faces(vert1: wp.array[wp.vec3], len1: int, vert2: wp.array[wp.vec3], len2: int) -> Tuple[int, wp.vec2i]: res = wp.vec2i() for i in range(len1): for j in range(len2): @@ -1417,9 +1415,7 @@ def _aligned_faces( # find two normals that are perpendicular to each other within a tolerance # return 1 if found @wp.func -def _aligned_face_edge( - edge: wp.array(dtype=wp.vec3), nedge: int, face: wp.array(dtype=wp.vec3), nface: int -) -> Tuple[int, wp.vec2i]: +def _aligned_face_edge(edge: wp.array[wp.vec3], nedge: int, face: wp.array[wp.vec3], nface: int) -> Tuple[int, wp.vec2i]: res = wp.vec2i() for i in range(nface): 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 @wp.func -def _intersect1( - a1: wp.array(dtype=int), a2: wp.array(dtype=int), start1: int, start2: int, len1: int, len2: int -) -> Tuple[int, wp.vec2i]: +def _intersect1(a1: wp.array[int], a2: wp.array[int], start1: int, start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]: count = int(0) res = wp.vec2i() for i in range(start1, start1 + len1): @@ -1448,7 +1442,7 @@ def _intersect1( @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) res = wp.vec2i() for i in range(len1): @@ -1470,13 +1464,13 @@ def _mesh_normals( mat: wp.mat33, vertadr: int, polyadr: int, - polynormal: wp.array(dtype=wp.vec3), - polymapadr: wp.array(dtype=int), - polymapnum: wp.array(dtype=int), - polymap: wp.array(dtype=int), + polynormal: wp.array[wp.vec3], + polymapadr: wp.array[int], + polymapnum: wp.array[int], + polymap: wp.array[int], # Out: - normals_out: wp.array(dtype=wp.vec3), - indices_out: wp.array(dtype=int), + normals_out: wp.array[wp.vec3], + indices_out: wp.array[int], ) -> int: v1 = feature_index[0] v2 = feature_index[1] @@ -1540,19 +1534,19 @@ def _mesh_edge_normals( pos: wp.vec3, vertadr: int, polyadr: int, - vert: wp.array(dtype=wp.vec3), - polyvertadr: wp.array(dtype=int), - polyvertnum: wp.array(dtype=int), - polyvert: wp.array(dtype=int), - polymapadr: wp.array(dtype=int), - polymapnum: wp.array(dtype=int), - polymap: wp.array(dtype=int), + vert: wp.array[wp.vec3], + polyvertadr: wp.array[int], + polyvertnum: wp.array[int], + polyvert: wp.array[int], + polymapadr: wp.array[int], + polymapnum: wp.array[int], + polymap: wp.array[int], v1: wp.vec3, v2: wp.vec3, v1i: int, # Out: - normals_out: wp.array(dtype=wp.vec3), - endverts_out: wp.array(dtype=wp.vec3), + normals_out: wp.array[wp.vec3], + endverts_out: wp.array[wp.vec3], ) -> int: # only one edge if dim == 2: @@ -1586,8 +1580,8 @@ def _box_normals2( mat: wp.mat33, n: wp.vec3, # Out: - normal_out: wp.array(dtype=wp.vec3), - index_out: wp.array(dtype=int), + normal_out: wp.array[wp.vec3], + index_out: wp.array[int], ) -> int: # 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) @@ -1621,8 +1615,8 @@ def _box_normals( mat: wp.mat33, dir: wp.vec3, # Out: - normal_out: wp.array(dtype=wp.vec3), - index_out: wp.array(dtype=int), + normal_out: wp.array[wp.vec3], + index_out: wp.array[int], ) -> int: v1 = feature_index[0] v2 = feature_index[1] @@ -1669,7 +1663,7 @@ def _box_normals( # c is 1 if edge is diagonal of a box face # c is 2 if edge is an external edge of box if c == 1 or c == 2: - return 2 + return c return _box_normals2(mat, dir, normal_out, index_out) if feature_dim == 1: @@ -1698,8 +1692,8 @@ def _box_edge_normals( v2: wp.vec3, v1i: int, # Out: - normal_out: wp.array(dtype=wp.vec3), - endvert_out: wp.array(dtype=wp.vec3), + normal_out: wp.array[wp.vec3], + endvert_out: wp.array[wp.vec3], ) -> int: if dim == 2: endvert_out[0] = v2 @@ -1726,7 +1720,7 @@ def _box_edge_normals( # recover face of a box from its index @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 if idx == 0: # right face_out[0] = mat @ wp.vec3(size[0], size[1], size[2]) + pos @@ -1775,13 +1769,13 @@ def _mesh_face( pos: wp.vec3, vertadr: int, polyadr: int, - vert: wp.array(dtype=wp.vec3), - polyvertadr: wp.array(dtype=int), - polyvertnum: wp.array(dtype=int), - polyvert: wp.array(dtype=int), + vert: wp.array[wp.vec3], + polyvertadr: wp.array[int], + polyvertnum: wp.array[int], + polyvert: wp.array[int], idx: int, # Out: - face_out: wp.array(dtype=wp.vec3), + face_out: wp.array[wp.vec3], ) -> int: adr = polyvertadr[polyadr + idx] j = int(0) @@ -1821,17 +1815,17 @@ def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> float: @wp.func def _polygon_clip( # In: - plane_normal: wp.array(dtype=wp.vec3), - plane_dist: wp.array(dtype=float), - face1: wp.array(dtype=wp.vec3), + plane_normal: wp.array[wp.vec3], + plane_dist: wp.array[float], + face1: wp.array[wp.vec3], nface1: int, - face2: wp.array(dtype=wp.vec3), + face2: wp.array[wp.vec3], nface2: int, n: wp.vec3, dir: wp.vec3, # Out: - polygon_out: wp.array(dtype=wp.vec3), - clipped_out: wp.array(dtype=wp.vec3), + polygon_out: wp.array[wp.vec3], + clipped_out: wp.array[wp.vec3], ) -> Tuple[int, mat43, mat43]: witness1 = mat43() witness2 = mat43() @@ -1918,13 +1912,13 @@ def _polygon_clip( @wp.func def _set_edge( # In: - vert1: wp.array(dtype=wp.vec3), - vert2: wp.array(dtype=wp.vec3), + vert1: wp.array[wp.vec3], + vert2: wp.array[wp.vec3], start: int, end: int, offset: int, # Out: - face_out: wp.array(dtype=wp.vec3), + face_out: wp.array[wp.vec3], ) -> int: face_out[0] = vert1[2 * start + offset] face_out[1] = vert2[end] @@ -1935,19 +1929,19 @@ def _set_edge( @wp.func def multicontact( # In: - polygon: wp.array(dtype=wp.vec3), - clipped: wp.array(dtype=wp.vec3), - plane_normal: wp.array(dtype=wp.vec3), - plane_dist: wp.array(dtype=float), - idx1: wp.array(dtype=int), - idx2: wp.array(dtype=int), - n1: wp.array(dtype=wp.vec3), - n2: wp.array(dtype=wp.vec3), - endvert: wp.array(dtype=wp.vec3), - face1: wp.array(dtype=wp.vec3), - face2: wp.array(dtype=wp.vec3), - epa_vert: wp.array(dtype=wp.vec3), - epa_vert_index: wp.array(dtype=int), + polygon: wp.array[wp.vec3], + clipped: wp.array[wp.vec3], + plane_normal: wp.array[wp.vec3], + plane_dist: wp.array[float], + idx1: wp.array[int], + idx2: wp.array[int], + n1: wp.array[wp.vec3], + n2: wp.array[wp.vec3], + endvert: wp.array[wp.vec3], + face1: wp.array[wp.vec3], + face2: wp.array[wp.vec3], + epa_vert: wp.array[wp.vec3], + epa_vert_index: wp.array[int], epa_face: int, x1: wp.vec3, x2: wp.vec3, @@ -2215,12 +2209,12 @@ def ccd( geomtype2: int, x_1: wp.vec3, x_2: wp.vec3, - vert: wp.array(dtype=wp.vec3), - vert_index: wp.array(dtype=int), - face: wp.array(dtype=int), - face_pr: wp.array(dtype=wp.vec3), - face_norm2: wp.array(dtype=float), - horizon: wp.array(dtype=int), + vert: wp.array[wp.vec3], + vert_index: wp.array[int], + face: wp.array[int], + face_pr: wp.array[wp.vec3], + face_norm2: wp.array[float], + horizon: wp.array[int], ) -> Tuple[float, int, wp.vec3, wp.vec3, int]: """General convex collision detection via GJK/EPA.""" full_margin1 = 0.0 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py index f1829de4..1002c727 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py @@ -295,21 +295,21 @@ def plane_sphere_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contact between a sphere and a plane.""" normal = plane.normal @@ -367,21 +367,21 @@ def sphere_sphere_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contact between two spheres.""" 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, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates one contact between a sphere and a capsule.""" # capsule axis @@ -512,21 +512,21 @@ def capsule_capsule_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between two capsules.""" # capsule axes @@ -598,21 +598,21 @@ def plane_capsule_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between a capsule and a plane.""" # capsule axis @@ -680,21 +680,21 @@ def plane_ellipsoid_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between an ellipsoid and a plane.""" 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, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between a box and a plane.""" 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, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between a plane and a convex object.""" dist, pos, normal = plane_convex(plane.normal, plane.pos, convex) @@ -897,21 +897,21 @@ def sphere_cylinder_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between a sphere and a cylinder.""" # cylinder axis @@ -978,21 +978,21 @@ def plane_cylinder_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between a cylinder and a plane.""" # cylinder axis @@ -1061,21 +1061,21 @@ def sphere_box_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): 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, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between a capsule and a box.""" # Extract capsule axis @@ -1216,21 +1216,21 @@ def box_box_wrapper( geoms: wp.vec2i, pairid: wp.vec2i, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): """Calculates contacts between two boxes.""" # 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) def primitive_narrowphase( # Model: - geom_type: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_graph: wp.array(dtype=int), - mesh_polynum: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), + geom_type: wp.array[int], + geom_condim: wp.array[int], + geom_dataid: wp.array2d[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_size: wp.array2d[wp.vec3], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_graph: wp.array[int], + mesh_polynum: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polymap: wp.array[int], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + pair_friction: wp.array2d[vec5], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], naconmax_in: int, - ncollision_in: wp.array(dtype=int), + ncollision_in: wp.array[int], # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), - collision_worldid_in: wp.array(dtype=int), + collision_pair_in: wp.array[wp.vec2i], + collision_pairid_in: wp.array[wp.vec2i], + collision_worldid_in: wp.array[int], # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): tid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py index 4decf33c..5937cf6e 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py @@ -54,9 +54,9 @@ class AABB: class VolumeData: center: wp.vec3 half_size: wp.vec3 - oct_aabb: wp.array2d(dtype=wp.vec3) - oct_child: wp.array(dtype=vec8i) - oct_coeff: wp.array(dtype=vec8) + oct_aabb: wp.array2d[wp.vec3] + oct_child: wp.array[vec8i] + oct_coeff: wp.array[vec8] root: int = 0 valid: bool = False @@ -64,10 +64,10 @@ class VolumeData: @wp.struct class MeshData: nmeshface: int - mesh_vertadr: wp.array(dtype=int) - mesh_vert: wp.array(dtype=wp.vec3) - mesh_faceadr: wp.array(dtype=int) - mesh_face: wp.array(dtype=wp.vec3i) + mesh_vertadr: wp.array[int] + mesh_vert: wp.array[wp.vec3] + mesh_faceadr: wp.array[int] + mesh_face: wp.array[wp.vec3i] data_id: int pos: wp.vec3 mat: wp.mat33 @@ -80,12 +80,12 @@ class MeshData: @wp.func def get_sdf_params( # Model: - oct_child: wp.array(dtype=vec8i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8), - mesh_octadr: wp.array(dtype=int), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=vec_pluginattr), + oct_child: wp.array[vec8i], + oct_aabb: wp.array2d[wp.vec3], + oct_coeff: wp.array[vec8], + mesh_octadr: wp.array[int], + plugin: wp.array[int], + plugin_attr: wp.array[vec_pluginattr], # In: g_type: int, 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 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]]: stack = root niter = int(100) @@ -665,75 +665,75 @@ def gradient_descent( def _sdf_narrowphase( # Model: nmeshface: int, - oct_child: wp.array(dtype=vec8i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8), - geom_type: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_faceadr: wp.array(dtype=int), - mesh_octadr: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_face: wp.array(dtype=wp.vec3i), - mesh_graph: wp.array(dtype=int), - mesh_polynum: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=vec_pluginattr), - geom_plugin_index: wp.array(dtype=int), + oct_child: wp.array[vec8i], + oct_aabb: wp.array2d[wp.vec3], + oct_coeff: wp.array[vec8], + geom_type: wp.array[int], + geom_condim: wp.array[int], + geom_dataid: wp.array2d[int], + geom_priority: wp.array[int], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_solimp: wp.array2d[vec5], + geom_size: wp.array2d[wp.vec3], + geom_aabb: wp.array3d[wp.vec3], + geom_friction: wp.array2d[wp.vec3], + geom_margin: wp.array2d[float], + geom_gap: wp.array2d[float], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_faceadr: wp.array[int], + mesh_octadr: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_face: wp.array[wp.vec3i], + mesh_graph: wp.array[int], + mesh_polynum: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polymap: wp.array[int], + pair_dim: wp.array[int], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + pair_solimp: wp.array2d[vec5], + pair_margin: wp.array2d[float], + pair_gap: wp.array2d[float], + pair_friction: wp.array2d[vec5], + plugin: wp.array[int], + plugin_attr: wp.array[vec_pluginattr], + geom_plugin_index: wp.array[int], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], naconmax_in: int, - ncollision_in: wp.array(dtype=int), + ncollision_in: wp.array[int], # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), - collision_worldid_in: wp.array(dtype=int), + collision_pair_in: wp.array[wp.vec2i], + collision_pairid_in: wp.array[wp.vec2i], + collision_worldid_in: wp.array[int], sdf_initpoints: int, sdf_iterations: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], + nacon_out: wp.array[int], ): i, contact_tid = wp.tid() if i >= sdf_initpoints: @@ -817,12 +817,32 @@ def _sdf_narrowphase( pos1 = geom1.pos rot1 = geom1.rot + dataid_setid = worldid % geom_dataid.shape[0] + 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( - 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 @@ -830,7 +850,7 @@ def _sdf_narrowphase( mesh_data1.mesh_vert = mesh_vert mesh_data1.mesh_faceadr = mesh_faceadr 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.mat = geom1.rot mesh_data1.size = geom1.size @@ -843,7 +863,7 @@ def _sdf_narrowphase( mesh_data2.mesh_vert = mesh_vert mesh_data2.mesh_faceadr = mesh_faceadr 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.mat = geom2.rot mesh_data2.size = geom2.size diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py index eec47583..788f1af5 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py @@ -32,12 +32,12 @@ wp.set_module_options({"enable_backward": False}) @wp.kernel def _zero_constraint_counts( # Data out: - ne_out: wp.array(dtype=int), - nf_out: wp.array(dtype=int), - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), + ne_out: wp.array[int], + nf_out: wp.array[int], + nl_out: wp.array[int], + nefc_out: wp.array[int], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid = wp.tid() @@ -68,14 +68,14 @@ def _efc_row( type: int, id: int, # Out: - type_out: wp.array2d(dtype=int), - id_out: wp.array2d(dtype=int), - pos_out: wp.array2d(dtype=float), - margin_out: wp.array2d(dtype=float), - D_out: wp.array2d(dtype=float), - vel_out: wp.array2d(dtype=float), - aref_out: wp.array2d(dtype=float), - frictionloss_out: wp.array2d(dtype=float), + type_out: wp.array2d[int], + id_out: wp.array2d[int], + pos_out: wp.array2d[float], + margin_out: wp.array2d[float], + D_out: wp.array2d[float], + vel_out: wp.array2d[float], + aref_out: wp.array2d[float], + frictionloss_out: wp.array2d[float], ): # calculate kbi timeconst = solref[0] @@ -126,52 +126,52 @@ def _equality_connect( # Model: nv: int, nsite: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_weldid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + body_invweight0: wp.array2d[wp.vec2], + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + site_bodyid: wp.array[int], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], + eq_solref: wp.array2d[wp.vec2], + eq_solimp: wp.array2d[vec5], + eq_data: wp.array2d[vec11], is_sparse: bool, - eq_connect_adr: wp.array(dtype=int), + eq_connect_adr: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + qvel_in: wp.array2d[float], + eq_active_in: wp.array2d[bool], + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], njmax_in: int, njmax_nnz_in: int, # Data out: - ne_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + ne_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): """Calculates constraint rows for connect equality constraints.""" worldid, eqconnectid = wp.tid() @@ -368,42 +368,42 @@ def _equality_connect( def _equality_joint( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - qpos0: wp.array2d(dtype=float), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_invweight0: wp.array2d(dtype=float), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), + qpos0: wp.array2d[float], + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], + dof_invweight0: wp.array2d[float], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_solref: wp.array2d[wp.vec2], + eq_solimp: wp.array2d[vec5], + eq_data: wp.array2d[vec11], is_sparse: bool, - eq_jnt_adr: wp.array(dtype=int), + eq_jnt_adr: wp.array[int], # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), + qpos_in: wp.array2d[float], + qvel_in: wp.array2d[float], + eq_active_in: wp.array2d[bool], njmax_in: int, njmax_nnz_in: int, # Data out: - ne_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + ne_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, eqjntid = wp.tid() eqid = eq_jnt_adr[eqjntid] @@ -499,44 +499,44 @@ def _equality_joint( def _equality_tendon( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_length0: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_solref: wp.array2d[wp.vec2], + eq_solimp: wp.array2d[vec5], + eq_data: wp.array2d[vec11], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_length0: wp.array2d[float], + tendon_invweight0: wp.array2d[float], is_sparse: bool, - eq_ten_adr: wp.array(dtype=int), + eq_ten_adr: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - ten_J_in: wp.array2d(dtype=float), - ten_length_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + eq_active_in: wp.array2d[bool], + ten_J_in: wp.array2d[float], + ten_length_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - ne_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + ne_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, eqtenid = wp.tid() eqid = eq_ten_adr[eqtenid] @@ -679,42 +679,42 @@ def _equality_flex(is_sparse: bool): def kernel( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - flex_edgeadr: wp.array(dtype=int), - flex_edgenum: wp.array(dtype=int), - flexedge_length0: wp.array(dtype=float), - flexedge_invweight0: wp.array(dtype=float), - flexedge_J_rownnz: wp.array(dtype=int), - flexedge_J_rowadr: wp.array(dtype=int), - flexedge_J_colind: wp.array(dtype=int), - eq_obj1id: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_flex_adr: wp.array(dtype=int), + flex_edgeadr: wp.array[int], + flex_edgenum: wp.array[int], + flexedge_length0: wp.array[float], + flexedge_invweight0: wp.array[float], + flexedge_J_rownnz: wp.array[int], + flexedge_J_rowadr: wp.array[int], + flexedge_J_colind: wp.array[int], + eq_obj1id: wp.array[int], + eq_solref: wp.array2d[wp.vec2], + eq_solimp: wp.array2d[vec5], + eq_flex_adr: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - flexedge_J_in: wp.array2d(dtype=float), - flexedge_length_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + flexedge_J_in: wp.array2d[float], + flexedge_length_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - ne_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + ne_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, eqflexid, edgeid = wp.tid() eqid = eq_flex_adr[eqflexid] @@ -794,54 +794,54 @@ def _equality_weld( # Model: nv: int, nsite: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_quat: wp.array2d(dtype=wp.quat), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_weldid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + body_invweight0: wp.array2d[wp.vec2], + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + site_bodyid: wp.array[int], + site_quat: wp.array2d[wp.quat], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], + eq_solref: wp.array2d[wp.vec2], + eq_solimp: wp.array2d[vec5], + eq_data: wp.array2d[vec11], is_sparse: bool, - eq_wld_adr: wp.array(dtype=int), + eq_wld_adr: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), - xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + qvel_in: wp.array2d[float], + eq_active_in: wp.array2d[bool], + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], + xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], njmax_in: int, njmax_nnz_in: int, # Data out: - ne_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + ne_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, eqweldid = wp.tid() eqid = eq_wld_adr[eqweldid] @@ -1114,34 +1114,34 @@ def _equality_weld( def _friction_dof( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - dof_solref: wp.array2d(dtype=wp.vec2), - dof_solimp: wp.array2d(dtype=vec5), - dof_frictionloss: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), + dof_solref: wp.array2d[wp.vec2], + dof_solimp: wp.array2d[vec5], + dof_frictionloss: wp.array2d[float], + dof_invweight0: wp.array2d[float], is_sparse: bool, # Data in: - qvel_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - nf_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nf_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, dofid = wp.tid() @@ -1204,38 +1204,38 @@ def _friction_dof( def _friction_tendon( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_solref_fri: wp.array2d(dtype=wp.vec2), - tendon_solimp_fri: wp.array2d(dtype=vec5), - tendon_frictionloss: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_solref_fri: wp.array2d[wp.vec2], + tendon_solimp_fri: wp.array2d[vec5], + tendon_frictionloss: wp.array2d[float], + tendon_invweight0: wp.array2d[float], is_sparse: bool, # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + ten_J_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - nf_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nf_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, tenid = wp.tid() @@ -1317,39 +1317,39 @@ def _friction_tendon( def _limit_slide_hinge( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array2d(dtype=vec5), - jnt_range: wp.array2d(dtype=wp.vec2), - jnt_margin: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_solref: wp.array2d[wp.vec2], + jnt_solimp: wp.array2d[vec5], + jnt_range: wp.array2d[wp.vec2], + jnt_margin: wp.array2d[float], + dof_invweight0: wp.array2d[float], is_sparse: bool, - jnt_limited_slide_hinge_adr: wp.array(dtype=int), + jnt_limited_slide_hinge_adr: wp.array[int], # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], + qvel_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nl_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, jntlimitedid = wp.tid() jntid = jnt_limited_slide_hinge_adr[jntlimitedid] @@ -1422,39 +1422,39 @@ def _limit_slide_hinge( def _limit_ball( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array2d(dtype=vec5), - jnt_range: wp.array2d(dtype=wp.vec2), - jnt_margin: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_solref: wp.array2d[wp.vec2], + jnt_solimp: wp.array2d[vec5], + jnt_range: wp.array2d[wp.vec2], + jnt_margin: wp.array2d[float], + dof_invweight0: wp.array2d[float], is_sparse: bool, - jnt_limited_ball_adr: wp.array(dtype=int), + jnt_limited_ball_adr: wp.array[int], # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], + qvel_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nl_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, jntlimitedid = wp.tid() jntid = jnt_limited_ball_adr[jntlimitedid] @@ -1547,41 +1547,41 @@ def _limit_ball( def _limit_tendon( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_solref_lim: wp.array2d(dtype=wp.vec2), - tendon_solimp_lim: wp.array2d(dtype=vec5), - tendon_range: wp.array2d(dtype=wp.vec2), - tendon_margin: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_solref_lim: wp.array2d[wp.vec2], + tendon_solimp_lim: wp.array2d[vec5], + tendon_range: wp.array2d[wp.vec2], + tendon_margin: wp.array2d[float], + tendon_invweight0: wp.array2d[float], is_sparse: bool, - tendon_limited_adr: wp.array(dtype=int), + tendon_limited_adr: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array2d(dtype=float), - ten_length_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + ten_J_in: wp.array2d[float], + ten_length_in: wp.array2d[float], njmax_in: int, njmax_nnz_in: int, # Data out: - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nl_out: wp.array[int], + nefc_out: wp.array[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): worldid, tenlimitedid = wp.tid() tenid = tendon_limited_adr[tenlimitedid] @@ -1669,59 +1669,59 @@ def _limit_tendon( def _contact_pyramidal( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - opt_impratio_invsqrt: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), + opt_impratio_invsqrt: wp.array[float], + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_weldid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + body_invweight0: wp.array2d[wp.vec2], + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + geom_bodyid: wp.array[int], + flex_vertadr: wp.array[int], + flex_vertbodyid: wp.array[int], is_sparse: bool, # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + qvel_in: wp.array2d[float], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], njmax_in: int, njmax_nnz_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - dist_in: wp.array(dtype=float), - condim_in: wp.array(dtype=int), - includemargin_in: wp.array(dtype=float), - worldid_in: wp.array(dtype=int), - geom_in: wp.array(dtype=wp.vec2i), - flex_in: wp.array(dtype=wp.vec2i), - vert_in: wp.array(dtype=wp.vec2i), - pos_in: wp.array(dtype=wp.vec3), - frame_in: wp.array(dtype=wp.mat33), - friction_in: wp.array(dtype=vec5), - solref_in: wp.array(dtype=wp.vec2), - solimp_in: wp.array(dtype=vec5), - type_in: wp.array(dtype=int), + dist_in: wp.array[float], + condim_in: wp.array[int], + includemargin_in: wp.array[float], + worldid_in: wp.array[int], + geom_in: wp.array[wp.vec2i], + flex_in: wp.array[wp.vec2i], + vert_in: wp.array[wp.vec2i], + pos_in: wp.array[wp.vec3], + frame_in: wp.array[wp.mat33], + friction_in: wp.array[vec5], + solref_in: wp.array[wp.vec2], + solimp_in: wp.array[vec5], + type_in: wp.array[int], # Data out: - nefc_out: wp.array(dtype=int), - contact_efc_address_out: wp.array2d(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nefc_out: wp.array[int], + contact_efc_address_out: wp.array2d[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): conid, dimid = wp.tid() @@ -1940,60 +1940,60 @@ def _contact_pyramidal( def _contact_elliptic( # Model: nv: int, - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - opt_impratio_invsqrt: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), + opt_impratio_invsqrt: wp.array[float], + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_weldid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + body_invweight0: wp.array2d[wp.vec2], + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + geom_bodyid: wp.array[int], + flex_vertadr: wp.array[int], + flex_vertbodyid: wp.array[int], is_sparse: bool, # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + qvel_in: wp.array2d[float], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], njmax_in: int, njmax_nnz_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - dist_in: wp.array(dtype=float), - condim_in: wp.array(dtype=int), - includemargin_in: wp.array(dtype=float), - worldid_in: wp.array(dtype=int), - geom_in: wp.array(dtype=wp.vec2i), - flex_in: wp.array(dtype=wp.vec2i), - vert_in: wp.array(dtype=wp.vec2i), - pos_in: wp.array(dtype=wp.vec3), - frame_in: wp.array(dtype=wp.mat33), - friction_in: wp.array(dtype=vec5), - solref_in: wp.array(dtype=wp.vec2), - solreffriction_in: wp.array(dtype=wp.vec2), - solimp_in: wp.array(dtype=vec5), - type_in: wp.array(dtype=int), + dist_in: wp.array[float], + condim_in: wp.array[int], + includemargin_in: wp.array[float], + worldid_in: wp.array[int], + geom_in: wp.array[wp.vec2i], + flex_in: wp.array[wp.vec2i], + vert_in: wp.array[wp.vec2i], + pos_in: wp.array[wp.vec3], + frame_in: wp.array[wp.mat33], + friction_in: wp.array[vec5], + solref_in: wp.array[wp.vec2], + solreffriction_in: wp.array[wp.vec2], + solimp_in: wp.array[vec5], + type_in: wp.array[int], # Data out: - nefc_out: wp.array(dtype=int), - contact_efc_address_out: wp.array2d(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_rownnz_out: wp.array2d(dtype=int), - efc_J_rowadr_out: wp.array2d(dtype=int), - efc_J_colind_out: wp.array3d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), + nefc_out: wp.array[int], + contact_efc_address_out: wp.array2d[int], + efc_type_out: wp.array2d[int], + efc_id_out: wp.array2d[int], + efc_J_rownnz_out: wp.array2d[int], + efc_J_rowadr_out: wp.array2d[int], + efc_J_colind_out: wp.array3d[int], + efc_J_out: wp.array3d[float], + efc_pos_out: wp.array2d[float], + efc_margin_out: wp.array2d[float], + efc_D_out: wp.array2d[float], + efc_vel_out: wp.array2d[float], + efc_aref_out: wp.array2d[float], + efc_frictionloss_out: wp.array2d[float], # Out: - efc_nnz_out: wp.array(dtype=int), + efc_nnz_out: wp.array[int], ): conid, dimid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py index 20da751a..bdcd81cb 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py @@ -31,27 +31,27 @@ wp.set_module_options({"enable_backward": False}) @wp.kernel def _qderiv_actuator_passive_vel( # Model: - opt_timestep: wp.array(dtype=float), - actuator_dyntype: wp.array(dtype=int), - actuator_gaintype: wp.array(dtype=int), - actuator_biastype: wp.array(dtype=int), - actuator_actadr: wp.array(dtype=int), - actuator_actnum: wp.array(dtype=int), - actuator_forcelimited: wp.array(dtype=bool), - actuator_actlimited: wp.array(dtype=bool), - actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_gainprm: wp.array2d(dtype=vec10f), - actuator_biasprm: wp.array2d(dtype=vec10f), - actuator_actearly: wp.array(dtype=bool), - actuator_forcerange: wp.array2d(dtype=wp.vec2), - actuator_actrange: wp.array2d(dtype=wp.vec2), + opt_timestep: wp.array[float], + actuator_dyntype: wp.array[int], + actuator_gaintype: wp.array[int], + actuator_biastype: wp.array[int], + actuator_actadr: wp.array[int], + actuator_actnum: wp.array[int], + actuator_forcelimited: wp.array[bool], + actuator_actlimited: wp.array[bool], + actuator_dynprm: wp.array2d[vec10f], + actuator_gainprm: wp.array2d[vec10f], + actuator_biasprm: wp.array2d[vec10f], + actuator_actearly: wp.array[bool], + actuator_forcerange: wp.array2d[wp.vec2], + actuator_actrange: wp.array2d[wp.vec2], # Data in: - act_in: wp.array2d(dtype=float), - ctrl_in: wp.array2d(dtype=float), - act_dot_in: wp.array2d(dtype=float), - actuator_force_in: wp.array2d(dtype=float), + act_in: wp.array2d[float], + ctrl_in: wp.array2d[float], + act_dot_in: wp.array2d[float], + actuator_force_in: wp.array2d[float], # Out: - vel_out: wp.array2d(dtype=float), + vel_out: wp.array2d[float], ): worldid, actid = wp.tid() @@ -121,16 +121,16 @@ def _qderiv_actuator_passive_actuation_dense( # Model: nu: int, # Data in: - moment_rownnz_in: wp.array2d(dtype=int), - moment_rowadr_in: wp.array2d(dtype=int), - moment_colind_in: wp.array2d(dtype=int), - actuator_moment_in: wp.array2d(dtype=float), + moment_rownnz_in: wp.array2d[int], + moment_rowadr_in: wp.array2d[int], + moment_colind_in: wp.array2d[int], + actuator_moment_in: wp.array2d[float], # In: - vel_in: wp.array2d(dtype=float), - qMi: wp.array(dtype=int), - qMj: wp.array(dtype=int), + vel_in: wp.array2d[float], + qMi: wp.array[int], + qMj: wp.array[int], # Out: - qDeriv_out: wp.array3d(dtype=float), + qDeriv_out: wp.array3d[float], ): worldid, elemid = wp.tid() @@ -171,18 +171,18 @@ def _qderiv_actuator_passive_actuation_dense( @wp.kernel def _qderiv_actuator_passive_actuation_sparse( # Model: - M_rownnz: wp.array(dtype=int), - M_rowadr: wp.array(dtype=int), + M_rownnz: wp.array[int], + M_rowadr: wp.array[int], # Data in: - moment_rownnz_in: wp.array2d(dtype=int), - moment_rowadr_in: wp.array2d(dtype=int), - moment_colind_in: wp.array2d(dtype=int), - actuator_moment_in: wp.array2d(dtype=float), + moment_rownnz_in: wp.array2d[int], + moment_rowadr_in: wp.array2d[int], + moment_colind_in: wp.array2d[int], + actuator_moment_in: wp.array2d[float], # In: - vel_in: wp.array2d(dtype=float), - qMj: wp.array(dtype=int), + vel_in: wp.array2d[float], + qMj: wp.array[int], # Out: - qDeriv_out: wp.array3d(dtype=float), + qDeriv_out: wp.array3d[float], ): worldid, actid = wp.tid() @@ -225,18 +225,18 @@ def _qderiv_actuator_passive_actuation_sparse( @wp.kernel def _qderiv_actuator_passive( # Model: - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], opt_disableflags: int, - dof_damping: wp.array2d(dtype=float), + dof_damping: wp.array2d[float], is_sparse: bool, # Data in: - qM_in: wp.array3d(dtype=float), + qM_in: wp.array3d[float], # In: - qMi: wp.array(dtype=int), - qMj: wp.array(dtype=int), - qDeriv_in: wp.array3d(dtype=float), + qMi: wp.array[int], + qMj: wp.array[int], + qDeriv_in: wp.array3d[float], # Out: - qDeriv_out: wp.array3d(dtype=float), + qDeriv_out: wp.array3d[float], ): worldid, elemid = wp.tid() @@ -267,19 +267,19 @@ def _qderiv_actuator_passive( def _qderiv_tendon_damping( # Model: ntendon: int, - opt_timestep: wp.array(dtype=float), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_damping: wp.array2d(dtype=float), + opt_timestep: wp.array[float], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_damping: wp.array2d[float], is_sparse: bool, # Data in: - ten_J_in: wp.array2d(dtype=float), + ten_J_in: wp.array2d[float], # In: - qMi: wp.array(dtype=int), - qMj: wp.array(dtype=int), + qMi: wp.array[int], + qMj: wp.array[int], # Out: - qDeriv_out: wp.array3d(dtype=float), + qDeriv_out: wp.array3d[float], ): worldid, elemid = wp.tid() dofiid = qMi[elemid] @@ -318,7 +318,7 @@ def _qderiv_tendon_damping( @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. Args: diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py index 64bdd91f..55c7f357 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -51,17 +51,17 @@ wp.set_module_options({"enable_backward": False}) @wp.kernel def _next_position( # Model: - opt_timestep: wp.array(dtype=float), - jnt_type: wp.array(dtype=int), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), + opt_timestep: wp.array[float], + jnt_type: wp.array[int], + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], + qvel_in: wp.array2d[float], # In: qvel_scale_in: float, # Data out: - qpos_out: wp.array2d(dtype=float), + qpos_out: wp.array2d[float], ): worldid, jntid = wp.tid() timestep = opt_timestep[worldid % opt_timestep.shape[0]] @@ -115,14 +115,14 @@ def _next_position( @wp.kernel def _next_velocity( # Model: - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], # Data in: - qvel_in: wp.array2d(dtype=float), - qacc_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + qacc_in: wp.array2d[float], # In: qacc_scale_in: float, # Data out: - qvel_out: wp.array2d(dtype=float), + qvel_out: wp.array2d[float], ): worldid, dofid = wp.tid() timestep = opt_timestep[worldid % opt_timestep.shape[0]] @@ -132,21 +132,21 @@ def _next_velocity( @wp.kernel def _next_activation( # Model: - opt_timestep: wp.array(dtype=float), - actuator_dyntype: wp.array(dtype=int), - actuator_actadr: wp.array(dtype=int), - actuator_actnum: wp.array(dtype=int), - actuator_actlimited: wp.array(dtype=bool), - actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_actrange: wp.array2d(dtype=wp.vec2), + opt_timestep: wp.array[float], + actuator_dyntype: wp.array[int], + actuator_actadr: wp.array[int], + actuator_actnum: wp.array[int], + actuator_actlimited: wp.array[bool], + actuator_dynprm: wp.array2d[vec10f], + actuator_actrange: wp.array2d[wp.vec2], # Data in: - act_in: wp.array2d(dtype=float), - act_dot_in: wp.array2d(dtype=float), + act_in: wp.array2d[float], + act_dot_in: wp.array2d[float], # In: act_dot_scale: float, limit: bool, # Data out: - act_out: wp.array2d(dtype=float), + act_out: wp.array2d[float], ): worldid, uid = wp.tid() opt_timestep_id = worldid % opt_timestep.shape[0] @@ -171,21 +171,21 @@ def _next_activation( @wp.kernel def _next_time( # Model: - opt_timestep: wp.array(dtype=float), + opt_timestep: wp.array[float], is_sparse: bool, # Data in: - nefc_in: wp.array(dtype=int), - time_in: wp.array(dtype=float), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), + nefc_in: wp.array[int], + time_in: wp.array[float], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], nworld_in: int, naconmax_in: int, njmax_in: int, njmax_nnz_in: int, - nacon_in: wp.array(dtype=int), - ncollision_in: wp.array(dtype=int), + nacon_in: wp.array[int], + ncollision_in: wp.array[int], # Data out: - time_out: wp.array(dtype=float), + time_out: wp.array[float], ): worldid = wp.tid() 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 def _euler_damp_qfrc_sparse( # Model: - opt_timestep: wp.array(dtype=float), - dof_Madr: wp.array(dtype=int), - dof_damping: wp.array2d(dtype=float), + opt_timestep: wp.array[float], + dof_Madr: wp.array[int], + dof_damping: wp.array2d[float], # Out: - qM_integration_out: wp.array3d(dtype=float), + qM_integration_out: wp.array3d[float], ): worldid, tid = wp.tid() 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) def euler_dense( # Model: - opt_timestep: wp.array(dtype=float), - dof_damping: wp.array2d(dtype=float), + opt_timestep: wp.array[float], + dof_damping: wp.array2d[float], # Data in: - qM_in: wp.array3d(dtype=float), - efc_Ma_in: wp.array2d(dtype=float), + qM_in: wp.array3d[float], + efc_Ma_in: wp.array2d[float], # In: - adr_in: wp.array(dtype=int), + adr_in: wp.array[int], # Data out: - qacc_out: wp.array2d(dtype=float), + qacc_out: wp.array2d[float], ): worldid, nodeid = wp.tid() timestep = opt_timestep[worldid % opt_timestep.shape[0]] @@ -358,8 +358,8 @@ def _rk_perturb_state( m: Model, d: Data, scale: float, - qpos_t0: wp.array2d(dtype=float), - qvel_t0: wp.array2d(dtype=float), + qpos_t0: wp.array2d[float], + qvel_t0: wp.array2d[float], act_t0: Optional[wp.array] = None, ): # position @@ -403,13 +403,13 @@ def _rk_perturb_state( @wp.kernel def _rk_accumulate_velocity_acceleration( # Data in: - qvel_in: wp.array2d(dtype=float), - qacc_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + qacc_in: wp.array2d[float], # In: scale: float, # Data out: - qvel_out: wp.array2d(dtype=float), - qacc_out: wp.array2d(dtype=float), + qvel_out: wp.array2d[float], + qacc_out: wp.array2d[float], ): worldid, dofid = wp.tid() qvel_out[worldid, dofid] += scale * qvel_in[worldid, dofid] @@ -419,11 +419,11 @@ def _rk_accumulate_velocity_acceleration( @wp.kernel def _rk_accumulate_activation_velocity( # Data in: - act_dot_in: wp.array2d(dtype=float), + act_dot_in: wp.array2d[float], # In: scale: float, # Data out: - act_dot_out: wp.array2d(dtype=float), + act_dot_out: wp.array2d[float], ): worldid, actid = wp.tid() act_dot_out[worldid, actid] += scale * act_dot_in[worldid, actid] @@ -433,8 +433,8 @@ def _rk_accumulate( m: Model, d: Data, scale: float, - qvel_rk: wp.array2d(dtype=float), - qacc_rk: wp.array2d(dtype=float), + qvel_rk: wp.array2d[float], + qacc_rk: wp.array2d[float], 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.""" @@ -540,13 +540,13 @@ def fwd_position(m: Model, d: Data, factorize: bool = True): @wp.kernel def _actuator_velocity( # Data in: - qvel_in: wp.array2d(dtype=float), - moment_rownnz_in: wp.array2d(dtype=int), - moment_rowadr_in: wp.array2d(dtype=int), - moment_colind_in: wp.array2d(dtype=int), - actuator_moment_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + moment_rownnz_in: wp.array2d[int], + moment_rowadr_in: wp.array2d[int], + moment_colind_in: wp.array2d[int], + actuator_moment_in: wp.array2d[float], # Data out: - actuator_velocity_out: wp.array2d(dtype=float), + actuator_velocity_out: wp.array2d[float], ): worldid, actid = wp.tid() @@ -565,14 +565,14 @@ def _actuator_velocity( @wp.kernel def _tendon_velocity( # Model: - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], + ten_J_in: wp.array2d[float], # Data out: - ten_velocity_out: wp.array2d(dtype=float), + ten_velocity_out: wp.array2d[float], ): worldid, tenid = wp.tid() @@ -617,34 +617,34 @@ def fwd_velocity(m: Model, d: Data): def _actuator_force( # Model: na: int, - opt_timestep: wp.array(dtype=float), - actuator_dyntype: wp.array(dtype=int), - actuator_gaintype: wp.array(dtype=int), - actuator_biastype: wp.array(dtype=int), - actuator_actadr: wp.array(dtype=int), - actuator_actnum: wp.array(dtype=int), - actuator_ctrllimited: wp.array(dtype=bool), - actuator_forcelimited: wp.array(dtype=bool), - actuator_actlimited: wp.array(dtype=bool), - actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_gainprm: wp.array2d(dtype=vec10f), - actuator_biasprm: wp.array2d(dtype=vec10f), - actuator_actearly: wp.array(dtype=bool), - actuator_ctrlrange: wp.array2d(dtype=wp.vec2), - actuator_forcerange: wp.array2d(dtype=wp.vec2), - actuator_actrange: wp.array2d(dtype=wp.vec2), - actuator_acc0: wp.array2d(dtype=float), - actuator_lengthrange: wp.array2d(dtype=wp.vec2), + opt_timestep: wp.array[float], + actuator_dyntype: wp.array[int], + actuator_gaintype: wp.array[int], + actuator_biastype: wp.array[int], + actuator_actadr: wp.array[int], + actuator_actnum: wp.array[int], + actuator_ctrllimited: wp.array[bool], + actuator_forcelimited: wp.array[bool], + actuator_actlimited: wp.array[bool], + actuator_dynprm: wp.array2d[vec10f], + actuator_gainprm: wp.array2d[vec10f], + actuator_biasprm: wp.array2d[vec10f], + actuator_actearly: wp.array[bool], + actuator_ctrlrange: wp.array2d[wp.vec2], + actuator_forcerange: wp.array2d[wp.vec2], + actuator_actrange: wp.array2d[wp.vec2], + actuator_acc0: wp.array2d[float], + actuator_lengthrange: wp.array2d[wp.vec2], # Data in: - act_in: wp.array2d(dtype=float), - ctrl_in: wp.array2d(dtype=float), - actuator_length_in: wp.array2d(dtype=float), - actuator_velocity_in: wp.array2d(dtype=float), + act_in: wp.array2d[float], + ctrl_in: wp.array2d[float], + actuator_length_in: wp.array2d[float], + actuator_velocity_in: wp.array2d[float], # In: dsbl_clampctrl: int, # Data out: - act_dot_out: wp.array2d(dtype=float), - actuator_force_out: wp.array2d(dtype=float), + act_dot_out: wp.array2d[float], + actuator_force_out: wp.array2d[float], ): worldid, uid = wp.tid() @@ -738,12 +738,12 @@ def _actuator_force( @wp.kernel def _tendon_actuator_force( # Model: - actuator_trntype: wp.array(dtype=int), - actuator_trnid: wp.array(dtype=wp.vec2i), + actuator_trntype: wp.array[int], + actuator_trnid: wp.array[wp.vec2i], # Data in: - actuator_force_in: wp.array2d(dtype=float), + actuator_force_in: wp.array2d[float], # Out: - ten_actfrc_out: wp.array2d(dtype=float), + ten_actfrc_out: wp.array2d[float], ): worldid, actid = wp.tid() @@ -756,14 +756,14 @@ def _tendon_actuator_force( @wp.kernel def _tendon_actuator_force_clamp( # Model: - tendon_actfrclimited: wp.array(dtype=bool), - tendon_actfrcrange: wp.array2d(dtype=wp.vec2), - actuator_trntype: wp.array(dtype=int), - actuator_trnid: wp.array(dtype=wp.vec2i), + tendon_actfrclimited: wp.array[bool], + tendon_actfrcrange: wp.array2d[wp.vec2], + actuator_trntype: wp.array[int], + actuator_trnid: wp.array[wp.vec2i], # In: - ten_actfrc_in: wp.array2d(dtype=float), + ten_actfrc_in: wp.array2d[float], # Data out: - actuator_force_out: wp.array2d(dtype=float), + actuator_force_out: wp.array2d[float], ): worldid, actid = wp.tid() @@ -782,13 +782,13 @@ def _tendon_actuator_force_clamp( @wp.kernel def _qfrc_actuator( # Data in: - moment_rownnz_in: wp.array2d(dtype=int), - moment_rowadr_in: wp.array2d(dtype=int), - moment_colind_in: wp.array2d(dtype=int), - actuator_moment_in: wp.array2d(dtype=float), - actuator_force_in: wp.array2d(dtype=float), + moment_rownnz_in: wp.array2d[int], + moment_rowadr_in: wp.array2d[int], + moment_colind_in: wp.array2d[int], + actuator_moment_in: wp.array2d[float], + actuator_force_in: wp.array2d[float], # Data out: - qfrc_actuator_out: wp.array2d(dtype=float), + qfrc_actuator_out: wp.array2d[float], ): worldid, actid = wp.tid() @@ -806,15 +806,15 @@ def _qfrc_actuator( def _qfrc_actuator_gravcomp_limits( # Model: ngravcomp: int, - jnt_actfrclimited: wp.array(dtype=bool), - jnt_actgravcomp: wp.array(dtype=int), - jnt_actfrcrange: wp.array2d(dtype=wp.vec2), - dof_jntid: wp.array(dtype=int), + jnt_actfrclimited: wp.array[bool], + jnt_actgravcomp: wp.array[int], + jnt_actfrcrange: wp.array2d[wp.vec2], + dof_jntid: wp.array[int], # Data in: - qfrc_gravcomp_in: wp.array2d(dtype=float), - qfrc_actuator_in: wp.array2d(dtype=float), + qfrc_gravcomp_in: wp.array2d[float], + qfrc_actuator_in: wp.array2d[float], # Data out: - qfrc_actuator_out: wp.array2d(dtype=float), + qfrc_actuator_out: wp.array2d[float], ): worldid, dofid = wp.tid() jntid = dof_jntid[dofid] @@ -930,12 +930,12 @@ def fwd_actuation(m: Model, d: Data): @wp.kernel def _qfrc_smooth( # Data in: - qfrc_applied_in: wp.array2d(dtype=float), - qfrc_bias_in: wp.array2d(dtype=float), - qfrc_passive_in: wp.array2d(dtype=float), - qfrc_actuator_in: wp.array2d(dtype=float), + qfrc_applied_in: wp.array2d[float], + qfrc_bias_in: wp.array2d[float], + qfrc_passive_in: wp.array2d[float], + qfrc_actuator_in: wp.array2d[float], # Data out: - qfrc_smooth_out: wp.array2d(dtype=float), + qfrc_smooth_out: wp.array2d[float], ): worldid, dofid = wp.tid() qfrc_smooth_out[worldid, dofid] = ( diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py index 9fb9242b..065afde0 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py @@ -34,12 +34,12 @@ wp.set_module_options({"enable_backward": False}) @wp.kernel def _qfrc_eulerdamp( # Model: - opt_timestep: wp.array(dtype=float), - dof_damping: wp.array2d(dtype=float), + opt_timestep: wp.array[float], + dof_damping: wp.array2d[float], # Data in: - qacc_in: wp.array2d(dtype=float), + qacc_in: wp.array2d[float], # Out: - qfrc_out: wp.array2d(dtype=float), + qfrc_out: wp.array2d[float], ): worldid, dofid = wp.tid() timestep = opt_timestep[worldid % opt_timestep.shape[0]] @@ -49,13 +49,13 @@ def _qfrc_eulerdamp( @wp.kernel def _qfrc_inverse( # Data in: - qfrc_bias_in: wp.array2d(dtype=float), - qfrc_passive_in: wp.array2d(dtype=float), - qfrc_constraint_in: wp.array2d(dtype=float), + qfrc_bias_in: wp.array2d[float], + qfrc_passive_in: wp.array2d[float], + qfrc_constraint_in: wp.array2d[float], # In: - Ma: wp.array2d(dtype=float), + Ma: wp.array2d[float], # Data out: - qfrc_inverse_out: wp.array2d(dtype=float), + qfrc_inverse_out: wp.array2d[float], ): worldid, dofid = wp.tid() @@ -67,7 +67,7 @@ def _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. Args: diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index 0b53094b..4a743681 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -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 -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. The array shape is determined by a field spec referencing MjModel / MjData array sizes. """ + spec_shape = getattr(spec, "shape", (0,)) shape = None - if spec.shape != (0,): - shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec.shape) + if spec_shape != (0,): + shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec_shape) if data is None and shape is None: return None # nothing to do @@ -50,7 +56,7 @@ def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array else: 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 array._is_batched = True # 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 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})) else: 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 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): - 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)) 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) - 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() 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 @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() 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, neq: int, nsensordata: int, - qpos0: wp.array2d(dtype=float), - eq_active0: wp.array(dtype=bool), + qpos0: wp.array2d[float], + eq_active0: wp.array[bool], # Data in: nworld_in: int, # In: - reset_in: wp.array(dtype=bool), + reset_in: wp.array[bool], # Data out: - solver_niter_out: wp.array(dtype=int), - ne_out: wp.array(dtype=int), - nf_out: wp.array(dtype=int), - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - time_out: wp.array(dtype=float), - energy_out: wp.array(dtype=wp.vec2), - qpos_out: wp.array2d(dtype=float), - qvel_out: wp.array2d(dtype=float), - act_out: wp.array2d(dtype=float), - qacc_warmstart_out: wp.array2d(dtype=float), - ctrl_out: wp.array2d(dtype=float), - qfrc_applied_out: wp.array2d(dtype=float), - eq_active_out: wp.array2d(dtype=bool), - qacc_out: wp.array2d(dtype=float), - act_dot_out: wp.array2d(dtype=float), - sensordata_out: wp.array2d(dtype=float), - nacon_out: wp.array(dtype=int), + solver_niter_out: wp.array[int], + ne_out: wp.array[int], + nf_out: wp.array[int], + nl_out: wp.array[int], + nefc_out: wp.array[int], + time_out: wp.array[float], + energy_out: wp.array[wp.vec2], + qpos_out: wp.array2d[float], + qvel_out: wp.array2d[float], + act_out: wp.array2d[float], + qacc_warmstart_out: wp.array2d[float], + ctrl_out: wp.array2d[float], + qfrc_applied_out: wp.array2d[float], + eq_active_out: wp.array2d[bool], + qacc_out: wp.array2d[float], + act_dot_out: wp.array2d[float], + sensordata_out: wp.array2d[float], + nacon_out: wp.array[int], ): 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) def reset_mocap( # Model: - body_mocapid: wp.array(dtype=int), - body_pos: wp.array2d(dtype=wp.vec3), - body_quat: wp.array2d(dtype=wp.quat), + body_mocapid: wp.array[int], + body_pos: wp.array2d[wp.vec3], + body_quat: wp.array2d[wp.quat], # In: - reset_in: wp.array(dtype=bool), + reset_in: wp.array[bool], # Data out: - mocap_pos_out: wp.array2d(dtype=wp.vec3), - mocap_quat_out: wp.array2d(dtype=wp.quat), + mocap_pos_out: wp.array2d[wp.vec3], + mocap_quat_out: wp.array2d[wp.quat], ): 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) def reset_contact( # Data in: - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - reset_in: wp.array(dtype=bool), + reset_in: wp.array[bool], nefcaddress: int, # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=types.vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=types.vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_flex_out: wp.array(dtype=wp.vec2i), - contact_vert_out: wp.array(dtype=wp.vec2i), - contact_efc_address_out: wp.array2d(dtype=int), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), + contact_dist_out: wp.array[float], + contact_pos_out: wp.array[wp.vec3], + contact_frame_out: wp.array[wp.mat33], + contact_includemargin_out: wp.array[float], + contact_friction_out: wp.array[types.vec5], + contact_solref_out: wp.array[wp.vec2], + contact_solreffriction_out: wp.array[wp.vec2], + contact_solimp_out: wp.array[types.vec5], + contact_dim_out: wp.array[int], + contact_geom_out: wp.array[wp.vec2i], + contact_flex_out: wp.array[wp.vec2i], + contact_vert_out: wp.array[wp.vec2i], + contact_efc_address_out: wp.array2d[int], + contact_worldid_out: wp.array[int], + contact_type_out: wp.array[int], + contact_geomcollisionid_out: wp.array[int], ): conid = wp.tid() @@ -1697,8 +1703,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): # kernel_analyzer: off @wp.kernel def _init_subtreemass( - body_mass_in: wp.array2d(dtype=float), - body_subtreemass_out: wp.array2d(dtype=float), + body_mass_in: wp.array2d[float], + body_subtreemass_out: wp.array2d[float], ): worldid, bodyid = wp.tid() body_mass_id = worldid % body_mass_in.shape[0] @@ -1708,9 +1714,9 @@ def _init_subtreemass( @wp.kernel def _accumulate_subtreemass( - body_parentid: wp.array(dtype=int), - body_subtreemass_io: wp.array2d(dtype=float), - body_tree_: wp.array(dtype=int), + body_parentid: wp.array[int], + body_subtreemass_io: wp.array2d[float], + body_tree_: wp.array[int], ): worldid, nodeid = wp.tid() body_subtreemass_id = worldid % body_subtreemass_io.shape[0] @@ -1722,8 +1728,8 @@ def _accumulate_subtreemass( @wp.kernel def _copy_qpos0_to_qpos( - qpos0: wp.array2d(dtype=float), - qpos_out: wp.array2d(dtype=float), + qpos0: wp.array2d[float], + qpos_out: wp.array2d[float], ): worldid, i = wp.tid() qpos0_id = worldid % qpos0.shape[0] @@ -1732,8 +1738,8 @@ def _copy_qpos0_to_qpos( @wp.kernel def _copy_tendon_length0( - ten_length_in: wp.array2d(dtype=float), - tendon_length0_out: wp.array2d(dtype=float), + ten_length_in: wp.array2d[float], + tendon_length0_out: wp.array2d[float], ): worldid, tenid = wp.tid() tendon_length0_id = worldid % tendon_length0_out.shape[0] @@ -1744,9 +1750,9 @@ def _copy_tendon_length0( def _compute_meaninertia( nv: int, is_sparse: bool, - dof_Madr_in: wp.array(dtype=int), - qM_in: wp.array3d(dtype=float), - meaninertia_out: wp.array(dtype=float), + dof_Madr_in: wp.array[int], + qM_in: wp.array3d[float], + meaninertia_out: wp.array[float], ): """Compute mean diagonal inertia from qM at qpos0.""" worldid = wp.tid() @@ -1771,7 +1777,7 @@ def _compute_meaninertia( @wp.kernel def _set_unit_vector( dofid_target: int, - unit_vec_out: wp.array2d(dtype=float), + unit_vec_out: wp.array2d[float], ): worldid = wp.tid() nv = unit_vec_out.shape[1] @@ -1785,8 +1791,8 @@ def _set_unit_vector( @wp.kernel def _extract_dof_A_diag( dofid: int, - result_vec_in: wp.array2d(dtype=float), - dof_A_diag_out: wp.array2d(dtype=float), + result_vec_in: wp.array2d[float], + dof_A_diag_out: wp.array2d[float], ): worldid = wp.tid() dof_A_diag_id = worldid % dof_A_diag_out.shape[0] @@ -1795,11 +1801,11 @@ def _extract_dof_A_diag( @wp.kernel def _finalize_dof_invweight0( - dof_jntid: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_A_diag_in: wp.array2d(dtype=float), - dof_invweight0_out: wp.array2d(dtype=float), + dof_jntid: wp.array[int], + jnt_type: wp.array[int], + jnt_dofadr: wp.array[int], + dof_A_diag_in: wp.array2d[float], + dof_invweight0_out: wp.array2d[float], ): worldid, dofid = wp.tid() dof_invweight0_id = worldid % dof_invweight0_out.shape[0] @@ -1842,15 +1848,15 @@ def _compute_body_jac_row( nv: int, bodyid_target: int, row_idx: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - subtree_com_in: wp.array2d(dtype=wp.vec3), - xipos_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - body_jac_row_out: wp.array2d(dtype=float), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_dofadr: wp.array[int], + body_dofnum: wp.array[int], + dof_parentid: wp.array[int], + subtree_com_in: wp.array2d[wp.vec3], + xipos_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + body_jac_row_out: wp.array2d[float], ): worldid = wp.tid() @@ -1901,9 +1907,9 @@ def _compute_body_A_diag_entry( nv: int, bodyid_target: int, row_idx: int, - body_jac_row_in: wp.array2d(dtype=float), - result_vec_in: wp.array2d(dtype=float), - body_A_diag_out: wp.array3d(dtype=float), + body_jac_row_in: wp.array2d[float], + result_vec_in: wp.array2d[float], + body_A_diag_out: wp.array3d[float], ): worldid = wp.tid() body_A_diag_id = worldid % body_A_diag_out.shape[0] @@ -1916,9 +1922,9 @@ def _compute_body_A_diag_entry( @wp.kernel def _finalize_body_invweight0( - body_weldid: wp.array(dtype=int), - body_A_diag_in: wp.array3d(dtype=float), - body_invweight0_out: wp.array2d(dtype=wp.vec2), + body_weldid: wp.array[int], + body_A_diag_in: wp.array3d[float], + body_invweight0_out: wp.array2d[wp.vec2], ): worldid, bodyid = wp.tid() body_invweight0_id = worldid % body_invweight0_out.shape[0] @@ -1953,11 +1959,11 @@ def _finalize_body_invweight0( @wp.kernel def _copy_tendon_jacobian( tenid_target: int, - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - ten_J_in: wp.array2d(dtype=float), - ten_J_vec_out: wp.array2d(dtype=float), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + ten_J_in: wp.array2d[float], + ten_J_vec_out: wp.array2d[float], ): worldid = wp.tid() nv = ten_J_in.shape[2] @@ -1971,15 +1977,15 @@ def _copy_tendon_jacobian( @wp.kernel def _compute_tendon_dot_product( # Model: - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], # In: tenid_target: int, - ten_J_in: wp.array2d(dtype=float), - result_vec_in: wp.array2d(dtype=float), + ten_J_in: wp.array2d[float], + result_vec_in: wp.array2d[float], # Out: - tendon_invweight0_out: wp.array2d(dtype=float), + tendon_invweight0_out: wp.array2d[float], ): worldid = wp.tid() tendon_invweight0_id = worldid % tendon_invweight0_out.shape[0] @@ -1997,15 +2003,15 @@ def _compute_tendon_dot_product( @wp.kernel def _compute_cam_pos0( - cam_bodyid: wp.array(dtype=int), - cam_targetbodyid: wp.array(dtype=int), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), - xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cam_pos0_out: wp.array2d(dtype=wp.vec3), - cam_poscom0_out: wp.array2d(dtype=wp.vec3), - cam_mat0_out: wp.array2d(dtype=wp.mat33), + cam_bodyid: wp.array[int], + cam_targetbodyid: wp.array[int], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], + xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cam_pos0_out: wp.array2d[wp.vec3], + cam_poscom0_out: wp.array2d[wp.vec3], + cam_mat0_out: wp.array2d[wp.mat33], ): worldid, camid = wp.tid() cam_pos0_id = worldid % cam_pos0_out.shape[0] @@ -2023,15 +2029,15 @@ def _compute_cam_pos0( @wp.kernel def _compute_light_pos0( - light_bodyid: wp.array(dtype=int), - light_targetbodyid: wp.array(dtype=int), - light_xpos_in: wp.array2d(dtype=wp.vec3), - light_xdir_in: wp.array2d(dtype=wp.vec3), - xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - light_pos0_out: wp.array2d(dtype=wp.vec3), - light_poscom0_out: wp.array2d(dtype=wp.vec3), - light_dir0_out: wp.array2d(dtype=wp.vec3), + light_bodyid: wp.array[int], + light_targetbodyid: wp.array[int], + light_xpos_in: wp.array2d[wp.vec3], + light_xdir_in: wp.array2d[wp.vec3], + xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + light_pos0_out: wp.array2d[wp.vec3], + light_poscom0_out: wp.array2d[wp.vec3], + light_dir0_out: wp.array2d[wp.vec3], ): worldid, lightid = wp.tid() light_pos0_id = worldid % light_pos0_out.shape[0] @@ -2050,11 +2056,11 @@ def _compute_light_pos0( @wp.kernel def _copy_actuator_moment( actid_target: int, - moment_rownnz_in: wp.array2d(dtype=int), - moment_rowadr_in: wp.array2d(dtype=int), - moment_colind_in: wp.array2d(dtype=int), - actuator_moment_in: wp.array2d(dtype=float), - act_moment_vec_out: wp.array2d(dtype=float), + moment_rownnz_in: wp.array2d[int], + moment_rowadr_in: wp.array2d[int], + moment_colind_in: wp.array2d[int], + actuator_moment_in: wp.array2d[float], + act_moment_vec_out: wp.array2d[float], ): worldid = wp.tid() nv = act_moment_vec_out.shape[1] @@ -2072,8 +2078,8 @@ def _copy_actuator_moment( def _compute_actuator_acc0( actid_target: int, nv: int, - result_vec_in: wp.array2d(dtype=float), - actuator_acc0_out: wp.array2d(dtype=float), + result_vec_in: wp.array2d[float], + actuator_acc0_out: wp.array2d[float], ): worldid = wp.tid() norm_sq = float(0.0) @@ -2084,11 +2090,11 @@ def _compute_actuator_acc0( @wp.kernel def _compute_dof_M0( - dof_bodyid: wp.array(dtype=int), - dof_armature: wp.array2d(dtype=float), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - crb_in: wp.array2d(dtype=vec10), - dof_M0_out: wp.array2d(dtype=float), + dof_bodyid: wp.array[int], + dof_armature: wp.array2d[float], + cdof_in: wp.array2d[wp.spatial_vector], + crb_in: wp.array2d[vec10], + dof_M0_out: wp.array2d[float], ): worldid, dofid = wp.tid() bodyid = dof_bodyid[dofid] @@ -2099,15 +2105,15 @@ def _compute_dof_M0( @wp.kernel def _resolve_dampratio( - actuator_biastype: wp.array(dtype=int), - actuator_gainprm: wp.array2d(dtype=types.vec10f), - moment_rownnz_in: wp.array2d(dtype=int), - moment_rowadr_in: wp.array2d(dtype=int), - moment_colind_in: wp.array2d(dtype=int), - actuator_moment_in: wp.array2d(dtype=float), - dof_M0_in: wp.array2d(dtype=float), + actuator_biastype: wp.array[int], + actuator_gainprm: wp.array2d[types.vec10f], + moment_rownnz_in: wp.array2d[int], + moment_rowadr_in: wp.array2d[int], + moment_colind_in: wp.array2d[int], + actuator_moment_in: wp.array2d[float], + dof_M0_in: wp.array2d[float], nv: int, - actuator_biasprm: wp.array2d(dtype=types.vec10f), + actuator_biasprm: wp.array2d[types.vec10f], ): worldid, actid = wp.tid() biastype = actuator_biastype[actid] @@ -2150,15 +2156,15 @@ def _resolve_dampratio( @wp.kernel def _set_length_range( - actuator_trntype: wp.array(dtype=int), - actuator_trnid: wp.array(dtype=wp.vec2i), - actuator_gear: wp.array2d(dtype=wp.spatial_vector), - jnt_limited: wp.array(dtype=int), - jnt_range: wp.array2d(dtype=wp.vec2), - tendon_limited: wp.array(dtype=int), - tendon_range: wp.array2d(dtype=wp.vec2), + actuator_trntype: wp.array[int], + actuator_trnid: wp.array[wp.vec2i], + actuator_gear: wp.array2d[wp.spatial_vector], + jnt_limited: wp.array[int], + jnt_range: wp.array2d[wp.vec2], + tendon_limited: wp.array[int], + tendon_range: wp.array2d[wp.vec2], ntendon: int, - actuator_lengthrange_out: wp.array2d(dtype=wp.vec2), + actuator_lengthrange_out: wp.array2d[wp.vec2], ): worldid, actid = wp.tid() 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, }, } - 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"} if not isinstance(overrides, dict): @@ -2626,7 +2638,7 @@ def _build_rays( intrinsic: wp.vec4, znear: float, # Out: - ray_out: wp.array(dtype=wp.vec3), + ray_out: wp.array[wp.vec3], ): xid, yid = wp.tid() 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(): constructor = "cubql" - # Mesh BVHs + # Mesh BVHs – build for all meshes so per-world variants are available nmesh = mjm.nmesh 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] mesh_registry = {} 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)] - for mid in used_mesh_id: + for mid in range(nmesh): mesh, half = bvh.build_mesh_bvh(mjm, mid, constructor=constructor) mesh_registry[mesh.id] = mesh mesh_bvh_id[mid] = mesh.id diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py index c021db3f..f7fb3a10 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py @@ -26,24 +26,24 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope def _tree_edges( # Model: nv: int, - body_treeid: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_treeid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - eq_type: wp.array(dtype=int), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), + body_treeid: wp.array[int], + jnt_dofadr: wp.array[int], + dof_treeid: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + eq_type: wp.array[int], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], # Data in: - nefc_in: wp.array(dtype=int), - contact_geom_in: wp.array(dtype=wp.vec2i), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_J_in: wp.array3d(dtype=float), + nefc_in: wp.array[int], + contact_geom_in: wp.array[wp.vec2i], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_J_in: wp.array3d[float], njmax_in: int, # Out: - tree_tree: wp.array3d(dtype=int), # kernel_analyzer: off + tree_tree: wp.array3d[int], # kernel_analyzer: off ): """Find tree edges.""" worldid, efcid = wp.tid() @@ -151,7 +151,7 @@ def _tree_edges( 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.""" tree_tree.zero_() wp.launch( @@ -184,14 +184,14 @@ def _flood_fill( # Model: ntree: int, # In: - tree_tree_in: wp.array3d(dtype=int), - labels_in: wp.array2d(dtype=int), - stack_in: wp.array2d(dtype=int), + tree_tree_in: wp.array3d[int], + labels_in: wp.array2d[int], + stack_in: wp.array2d[int], # Data out: - nisland_out: wp.array(dtype=int), - tree_island_out: wp.array2d(dtype=int), + nisland_out: wp.array[int], + tree_island_out: wp.array2d[int], # Out: - stack_out: wp.array2d(dtype=int), + stack_out: wp.array2d[int], ): """DFS flood fill to discover islands using tree_tree matrix.""" worldid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py index 4abcff26..4c76d1f8 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -71,18 +71,18 @@ def _ellipsoid_max_moment(size: wp.vec3, dir: int) -> float: def _spring_damper_dof_passive( # Model: opt_disableflags: int, - qpos_spring: wp.array2d(dtype=float), - jnt_type: wp.array(dtype=int), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_stiffness: wp.array2d(dtype=float), - dof_damping: wp.array2d(dtype=float), + qpos_spring: wp.array2d[float], + jnt_type: wp.array[int], + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_stiffness: wp.array2d[float], + dof_damping: wp.array2d[float], # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], + qvel_in: wp.array2d[float], # Data out: - qfrc_spring_out: wp.array2d(dtype=float), - qfrc_damper_out: wp.array2d(dtype=float), + qfrc_spring_out: wp.array2d[float], + qfrc_damper_out: wp.array2d[float], ): worldid, jntid = wp.tid() dofid = jnt_dofadr[jntid] @@ -182,22 +182,22 @@ def _spring_damper_dof_passive( @wp.kernel def _spring_damper_tendon_passive( # Model: - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_stiffness: wp.array2d(dtype=float), - tendon_damping: wp.array2d(dtype=float), - tendon_lengthspring: wp.array2d(dtype=wp.vec2), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_stiffness: wp.array2d[float], + tendon_damping: wp.array2d[float], + tendon_lengthspring: wp.array2d[wp.vec2], # Data in: - ten_J_in: wp.array2d(dtype=float), - ten_length_in: wp.array2d(dtype=float), - ten_velocity_in: wp.array2d(dtype=float), + ten_J_in: wp.array2d[float], + ten_length_in: wp.array2d[float], + ten_velocity_in: wp.array2d[float], # In: dsbl_spring: bool, dsbl_damper: bool, # Data out: - qfrc_spring_out: wp.array2d(dtype=float), - qfrc_damper_out: wp.array2d(dtype=float), + qfrc_spring_out: wp.array2d[float], + qfrc_damper_out: wp.array2d[float], ): worldid, tenid, dofid_sparse = wp.tid() @@ -246,18 +246,18 @@ def _spring_damper_tendon_passive( @wp.kernel def _gravity_force( # Model: - opt_gravity: wp.array(dtype=wp.vec3), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_gravcomp: wp.array2d(dtype=float), - dof_bodyid: wp.array(dtype=int), + opt_gravity: wp.array[wp.vec3], + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_mass: wp.array2d[float], + body_gravcomp: wp.array2d[float], + dof_bodyid: wp.array[int], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + xipos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], # Data out: - qfrc_gravcomp_out: wp.array2d(dtype=float), + qfrc_gravcomp_out: wp.array2d[float], ): worldid, bodyid, dofid = wp.tid() bodyid += 1 # skip world body @@ -275,27 +275,27 @@ def _gravity_force( @wp.kernel def _fluid_force( # Model: - opt_wind: wp.array(dtype=wp.vec3), - opt_density: wp.array(dtype=float), - opt_viscosity: wp.array(dtype=float), - body_rootid: wp.array(dtype=int), - body_geomnum: wp.array(dtype=int), - body_geomadr: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_inertia: wp.array2d(dtype=wp.vec3), - geom_type: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_fluid: wp.array2d(dtype=float), - body_fluid_ellipsoid: wp.array(dtype=bool), + opt_wind: wp.array[wp.vec3], + opt_density: wp.array[float], + opt_viscosity: wp.array[float], + body_rootid: wp.array[int], + body_geomnum: wp.array[int], + body_geomadr: wp.array[int], + body_mass: wp.array2d[float], + body_inertia: wp.array2d[wp.vec3], + geom_type: wp.array[int], + geom_size: wp.array2d[wp.vec3], + geom_fluid: wp.array2d[float], + body_fluid_ellipsoid: wp.array[bool], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # 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.""" worldid, bodyid = wp.tid() @@ -535,18 +535,18 @@ def _fluid(m: Model, d: Data): @wp.kernel def _qfrc_passive( # Model: - jnt_actgravcomp: wp.array(dtype=int), - dof_jntid: wp.array(dtype=int), + jnt_actgravcomp: wp.array[int], + dof_jntid: wp.array[int], has_fluid: bool, # Data in: - qfrc_spring_in: wp.array2d(dtype=float), - qfrc_damper_in: wp.array2d(dtype=float), - qfrc_gravcomp_in: wp.array2d(dtype=float), - qfrc_fluid_in: wp.array2d(dtype=float), + qfrc_spring_in: wp.array2d[float], + qfrc_damper_in: wp.array2d[float], + qfrc_gravcomp_in: wp.array2d[float], + qfrc_fluid_in: wp.array2d[float], # In: gravcomp: bool, # Data out: - qfrc_passive_out: wp.array2d(dtype=float), + qfrc_passive_out: wp.array2d[float], ): worldid, dofid = wp.tid() qfrc_passive = qfrc_spring_in[worldid, dofid] @@ -567,29 +567,29 @@ def _qfrc_passive( def _flex_elasticity( # Model: nflex: int, - opt_timestep: wp.array(dtype=float), - body_dofadr: wp.array(dtype=int), - flex_dim: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_edgeadr: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elemedgeadr: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), - flex_elemedge: wp.array(dtype=int), - flexedge_length0: wp.array(dtype=float), - flex_stiffness: wp.array2d(dtype=float), - flex_damping: wp.array(dtype=float), + opt_timestep: wp.array[float], + body_dofadr: wp.array[int], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_edgeadr: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elemedgeadr: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_elem: wp.array[int], + flex_elemedge: wp.array[int], + flexedge_length0: wp.array[float], + flex_stiffness: wp.array2d[float], + flex_damping: wp.array[float], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), - flexedge_length_in: wp.array2d(dtype=float), - flexedge_velocity_in: wp.array2d(dtype=float), + flexvert_xpos_in: wp.array2d[wp.vec3], + flexedge_length_in: wp.array2d[float], + flexedge_velocity_in: wp.array2d[float], # In: dsbl_damper: bool, # Data out: - qfrc_spring_out: wp.array2d(dtype=float), + qfrc_spring_out: wp.array2d[float], ): worldid, elemid = wp.tid() timestep = opt_timestep[worldid % opt_timestep.shape[0]] @@ -665,19 +665,19 @@ def _flex_elasticity( def _flex_bending( # Model: nflex: int, - body_dofadr: wp.array(dtype=int), - flex_dim: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_edgeadr: wp.array(dtype=int), - flex_edgenum: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_edgeflap: wp.array(dtype=wp.vec2i), - flex_bending: wp.array2d(dtype=float), + body_dofadr: wp.array[int], + flex_dim: wp.array[int], + flex_vertadr: wp.array[int], + flex_edgeadr: wp.array[int], + flex_edgenum: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_edgeflap: wp.array[wp.vec2i], + flex_bending: wp.array2d[float], # Data in: - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + flexvert_xpos_in: wp.array2d[wp.vec3], # Data out: - qfrc_spring_out: wp.array2d(dtype=float), + qfrc_spring_out: wp.array2d[float], ): worldid, edgeid = wp.tid() nvert = 4 diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py index 57d6a5ba..a320644a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py @@ -52,12 +52,12 @@ def _ray_map(pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w @wp.func def _ray_eliminate( # Model: - body_weldid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_matid: wp.array(dtype=int), # kernel_analyzer: ignore - geom_group: wp.array(dtype=int), - geom_rgba: wp.array(dtype=wp.vec4), # kernel_analyzer: ignore - mat_rgba: wp.array(dtype=wp.vec4), # kernel_analyzer: ignore + body_weldid: wp.array[int], + geom_bodyid: wp.array[int], + geom_matid: wp.array[int], # kernel_analyzer: ignore + geom_group: wp.array[int], + geom_rgba: wp.array[wp.vec4], # kernel_analyzer: ignore + mat_rgba: wp.array[wp.vec4], # kernel_analyzer: ignore # In: geomid: int, 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 def ray_hfield( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - hfield_size: wp.array(dtype=wp.vec4), - hfield_nrow: wp.array(dtype=int), - hfield_ncol: wp.array(dtype=int), - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + hfield_size: wp.array[wp.vec4], + hfield_nrow: wp.array[int], + hfield_ncol: wp.array[int], + hfield_adr: wp.array[int], + hfield_data: wp.array[float], # In: pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3, id: int, + worldid: int, ) -> Tuple[float, wp.vec3]: # check geom type if geom_type[id] != GeomType.HFIELD: return -1.0, wp.vec3() # hfield id and dimensions - hid = geom_dataid[id] + hid = geom_dataid[worldid % geom_dataid.shape[0], id] nrow = hfield_nrow[hid] ncol = hfield_ncol[hid] @@ -622,10 +623,10 @@ def ray_hfield( def ray_mesh( # Model: nmeshface: int, - mesh_vertadr: wp.array(dtype=int), - mesh_faceadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_face: wp.array(dtype=wp.vec3i), + mesh_vertadr: wp.array[int], + mesh_faceadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_face: wp.array[wp.vec3i], # In: data_id: int, pos: wp.vec3, @@ -699,7 +700,7 @@ def ray_mesh( @wp.func def ray_mesh_with_bvh( # In: - mesh_bvh_id: wp.array(dtype=wp.uint64), + mesh_bvh_id: wp.array[wp.uint64], mesh_geom_id: int, pos: wp.vec3, mat: wp.mat33, @@ -732,7 +733,7 @@ def ray_mesh_with_bvh( @wp.func def ray_mesh_with_bvh_anyhit( # In: - mesh_bvh_id: wp.array(dtype=wp.uint64), + mesh_bvh_id: wp.array[wp.uint64], mesh_geom_id: int, pos: wp.vec3, mat: wp.mat33, @@ -752,7 +753,7 @@ def ray_mesh_with_bvh_anyhit( @wp.func def ray_flex_with_bvh( # In: - flex_bvh_id: wp.array(dtype=wp.uint64), + flex_bvh_id: wp.array[wp.uint64], flexid: int, group_root: int, pnt: wp.vec3, @@ -781,7 +782,7 @@ def ray_flex_with_bvh( @wp.func def ray_flex_with_bvh_anyhit( # In: - flex_bvh_id: wp.array(dtype=wp.uint64), + flex_bvh_id: wp.array[wp.uint64], flexid: int, group_root: int, 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( # Model: nmeshface: int, - body_weldid: wp.array(dtype=int), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_matid: wp.array2d(dtype=int), - geom_group: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array2d(dtype=wp.vec4), - mesh_vertadr: wp.array(dtype=int), - mesh_faceadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_face: wp.array(dtype=wp.vec3i), - hfield_size: wp.array(dtype=wp.vec4), - hfield_nrow: wp.array(dtype=int), - hfield_ncol: wp.array(dtype=int), - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), - mat_rgba: wp.array2d(dtype=wp.vec4), + body_weldid: wp.array[int], + geom_type: wp.array[int], + geom_bodyid: wp.array[int], + geom_dataid: wp.array2d[int], + geom_matid: wp.array2d[int], + geom_group: wp.array[int], + geom_size: wp.array2d[wp.vec3], + geom_rgba: wp.array2d[wp.vec4], + mesh_vertadr: wp.array[int], + mesh_faceadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_face: wp.array[wp.vec3i], + hfield_size: wp.array[wp.vec4], + hfield_nrow: wp.array[int], + hfield_ncol: wp.array[int], + hfield_adr: wp.array[int], + hfield_data: wp.array[float], + mat_rgba: wp.array2d[wp.vec4], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, pnt: wp.vec3, @@ -876,7 +877,7 @@ def _ray_geom_mesh( mesh_faceadr, mesh_vert, mesh_face, - geom_dataid[geomid], + geom_dataid[worldid % geom_dataid.shape[0], geomid], pos, mat, geom_size[worldid % geom_size.shape[0], geomid], @@ -897,6 +898,7 @@ def _ray_geom_mesh( pnt, vec, geomid, + worldid, ) else: return ray_geom(pos, mat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type) @@ -909,37 +911,37 @@ def _ray( # Model: ngeom: int, nmeshface: int, - body_weldid: wp.array(dtype=int), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_matid: wp.array2d(dtype=int), - geom_group: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array2d(dtype=wp.vec4), - mesh_vertadr: wp.array(dtype=int), - mesh_faceadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_face: wp.array(dtype=wp.vec3i), - hfield_size: wp.array(dtype=wp.vec4), - hfield_nrow: wp.array(dtype=int), - hfield_ncol: wp.array(dtype=int), - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), - mat_rgba: wp.array2d(dtype=wp.vec4), + body_weldid: wp.array[int], + geom_type: wp.array[int], + geom_bodyid: wp.array[int], + geom_dataid: wp.array2d[int], + geom_matid: wp.array2d[int], + geom_group: wp.array[int], + geom_size: wp.array2d[wp.vec3], + geom_rgba: wp.array2d[wp.vec4], + mesh_vertadr: wp.array[int], + mesh_faceadr: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_face: wp.array[wp.vec3i], + hfield_size: wp.array[wp.vec4], + hfield_nrow: wp.array[int], + hfield_ncol: wp.array[int], + hfield_adr: wp.array[int], + hfield_data: wp.array[float], + mat_rgba: wp.array2d[wp.vec4], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: - pnt: wp.array2d(dtype=wp.vec3), - vec: wp.array2d(dtype=wp.vec3), + pnt: wp.array2d[wp.vec3], + vec: wp.array2d[wp.vec3], geomgroup: vec6, flg_static: bool, - bodyexclude: wp.array(dtype=int), + bodyexclude: wp.array[int], # Out: - dist_out: wp.array2d(dtype=float), - geomid_out: wp.array2d(dtype=int), - normal_out: wp.array2d(dtype=wp.vec3), + dist_out: wp.array2d[float], + geomid_out: wp.array2d[int], + normal_out: wp.array2d[wp.vec3], ): worldid, rayid, tid = wp.tid() @@ -1011,18 +1013,18 @@ def _ray( @wp.func def _ray_geom_mesh_bvh( # Model: - body_weldid: wp.array(dtype=int), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_matid: wp.array2d(dtype=int), - geom_group: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array2d(dtype=wp.vec4), - mat_rgba: wp.array2d(dtype=wp.vec4), + body_weldid: wp.array[int], + geom_type: wp.array[int], + geom_bodyid: wp.array[int], + geom_dataid: wp.array2d[int], + geom_matid: wp.array2d[int], + geom_group: wp.array[int], + geom_size: wp.array2d[wp.vec3], + geom_rgba: wp.array2d[wp.vec4], + mat_rgba: wp.array2d[wp.vec4], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, pnt: wp.vec3, @@ -1031,8 +1033,8 @@ def _ray_geom_mesh_bvh( flg_static: bool, bodyexclude: int, geomid: int, - mesh_bvh_id: wp.array(dtype=wp.uint64), - hfield_bvh_id: wp.array(dtype=wp.uint64), + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], min_dist: float, ) -> Tuple[float, wp.vec3]: 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 t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh( bvh_ids, - geom_dataid[geomid], + geom_dataid[worldid % geom_dataid.shape[0], geomid], pos, mat, pnt, @@ -1081,33 +1083,33 @@ def _ray_geom_mesh_bvh( def _ray_bvh( # Model: ngeom: int, - body_weldid: wp.array(dtype=int), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_matid: wp.array2d(dtype=int), - geom_group: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array2d(dtype=wp.vec4), - mat_rgba: wp.array2d(dtype=wp.vec4), + body_weldid: wp.array[int], + geom_type: wp.array[int], + geom_bodyid: wp.array[int], + geom_dataid: wp.array2d[int], + geom_matid: wp.array2d[int], + geom_group: wp.array[int], + geom_size: wp.array2d[wp.vec3], + geom_rgba: wp.array2d[wp.vec4], + mat_rgba: wp.array2d[wp.vec4], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], # In: - pnt: wp.array2d(dtype=wp.vec3), - vec: wp.array2d(dtype=wp.vec3), + pnt: wp.array2d[wp.vec3], + vec: wp.array2d[wp.vec3], geomgroup: vec6, flg_static: bool, - bodyexclude: wp.array(dtype=int), + bodyexclude: wp.array[int], bvh_id: wp.uint64, - group_root: wp.array(dtype=int), - enabled_geom_ids: wp.array(dtype=int), - mesh_bvh_id: wp.array(dtype=wp.uint64), - hfield_bvh_id: wp.array(dtype=wp.uint64), + group_root: wp.array[int], + enabled_geom_ids: wp.array[int], + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], # Out: - dist_out: wp.array2d(dtype=float), - geomid_out: wp.array2d(dtype=int), - normal_out: wp.array2d(dtype=wp.vec3), + dist_out: wp.array2d[float], + geomid_out: wp.array2d[int], + normal_out: wp.array2d[wp.vec3], ): worldid, rayid = wp.tid() @@ -1166,8 +1168,8 @@ def _ray_bvh( def ray( m: Model, d: Data, - pnt: wp.array2d(dtype=wp.vec3), - vec: wp.array2d(dtype=wp.vec3), + pnt: wp.array2d[wp.vec3], + vec: wp.array2d[wp.vec3], geomgroup: vec6 | None = None, flg_static: bool = True, bodyexclude: int = -1, @@ -1210,14 +1212,14 @@ def ray( def rays( m: Model, d: Data, - pnt: wp.array2d(dtype=wp.vec3), - vec: wp.array2d(dtype=wp.vec3), + pnt: wp.array2d[wp.vec3], + vec: wp.array2d[wp.vec3], geomgroup: vec6, flg_static: bool, - bodyexclude: wp.array(dtype=int), - dist: wp.array2d(dtype=float), - geomid: wp.array2d(dtype=int), - normal: wp.array2d(dtype=wp.vec3), + bodyexclude: wp.array[int], + dist: wp.array2d[float], + geomid: wp.array2d[int], + normal: wp.array2d[wp.vec3], rc: RenderContext | None = None, ): """Ray intersection for multiple worlds and multiple rays. diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py index bc8d16c3..371032e1 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py @@ -43,17 +43,17 @@ wp.set_module_options({"enable_backward": False}) @wp.func def sample_texture( # Model: - geom_type: wp.array(dtype=int), - mesh_faceadr: wp.array(dtype=int), + geom_type: wp.array[int], + mesh_faceadr: wp.array[int], # In: geom_id: int, tex_repeat: wp.vec2, tex: wp.Texture2D, pos: wp.vec3, rot: wp.mat33, - mesh_facetexcoord: wp.array(dtype=wp.vec3i), - mesh_texcoord: wp.array(dtype=wp.vec2), - mesh_texcoord_offsets: wp.array(dtype=int), + mesh_facetexcoord: wp.array[wp.vec3i], + mesh_texcoord: wp.array[wp.vec2], + mesh_texcoord_offsets: wp.array[int], hit_point: wp.vec3, bary_u: float, bary_v: float, @@ -88,29 +88,29 @@ def sample_texture( @wp.func def cast_ray( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - flex_vertadr: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_radius: wp.array(dtype=float), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + flex_vertadr: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], # In: bvh_id: wp.uint64, group_root: int, worldid: int, bvh_ngeom: int, flex_bvh_ngeom: int, - enabled_geom_ids: wp.array(dtype=int), - mesh_bvh_id: wp.array(dtype=wp.uint64), - hfield_bvh_id: wp.array(dtype=wp.uint64), - flex_geom_flexid: wp.array(dtype=int), - flex_geom_edgeid: wp.array(dtype=int), - flex_bvh_id: wp.array(dtype=wp.uint64), - flex_group_root: wp.array2d(dtype=int), + enabled_geom_ids: wp.array[int], + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], + flex_bvh_id: wp.array[wp.uint64], + flex_group_root: wp.array2d[int], ray_origin_world: wp.vec3, ray_dir_world: wp.vec3, ) -> Tuple[int, float, wp.vec3, float, float, int, int]: @@ -159,7 +159,7 @@ def cast_ray( if gtype == GeomType.HFIELD: d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( hfield_bvh_id, - geom_dataid[gi], + geom_dataid[worldid % geom_dataid.shape[0], gi], geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], ray_origin_world, @@ -208,7 +208,7 @@ def cast_ray( if gtype == GeomType.MESH: d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh( mesh_bvh_id, - geom_dataid[gi], + geom_dataid[worldid % geom_dataid.shape[0], gi], geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], ray_origin_world, @@ -256,29 +256,29 @@ def cast_ray( @wp.func def cast_ray_first_hit( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - flex_vertadr: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_radius: wp.array(dtype=float), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + flex_vertadr: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], # In: bvh_id: wp.uint64, group_root: int, worldid: int, bvh_ngeom: int, bvh_nflexgeom: int, - enabled_geom_ids: wp.array(dtype=int), - mesh_bvh_id: wp.array(dtype=wp.uint64), - hfield_bvh_id: wp.array(dtype=wp.uint64), - flex_geom_flexid: wp.array(dtype=int), - flex_geom_edgeid: wp.array(dtype=int), - flex_bvh_id: wp.array(dtype=wp.uint64), - flex_group_root: wp.array2d(dtype=int), + enabled_geom_ids: wp.array[int], + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], + flex_bvh_id: wp.array[wp.uint64], + flex_group_root: wp.array2d[int], ray_origin_world: wp.vec3, ray_dir_world: wp.vec3, max_dist: float, @@ -314,7 +314,7 @@ def cast_ray_first_hit( if gtype == GeomType.HFIELD: d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( hfield_bvh_id, - geom_dataid[gi], + geom_dataid[worldid % geom_dataid.shape[0], gi], geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], ray_origin_world, @@ -363,7 +363,7 @@ def cast_ray_first_hit( if gtype == GeomType.MESH: hit = ray_mesh_with_bvh_anyhit( mesh_bvh_id, - geom_dataid[gi], + geom_dataid[worldid % geom_dataid.shape[0], gi], geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], ray_origin_world, @@ -409,30 +409,30 @@ def cast_ray_first_hit( @wp.func def compute_lighting( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - flex_vertadr: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_radius: wp.array(dtype=float), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + flex_vertadr: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], # In: use_shadows: bool, bvh_id: wp.uint64, group_root: int, bvh_ngeom: int, bvh_nflexgeom: int, - enabled_geom_ids: wp.array(dtype=int), + enabled_geom_ids: wp.array[int], worldid: int, - mesh_bvh_id: wp.array(dtype=wp.uint64), - hfield_bvh_id: wp.array(dtype=wp.uint64), - flex_geom_flexid: wp.array(dtype=int), - flex_geom_edgeid: wp.array(dtype=int), - flex_bvh_id: wp.array(dtype=wp.uint64), - flex_group_root: wp.array2d(dtype=int), + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], + flex_bvh_id: wp.array[wp.uint64], + flex_group_root: wp.array2d[int], lightactive: bool, lighttype: int, lightcastshadow: bool, @@ -530,65 +530,65 @@ def render(m: Model, d: Data, rc: RenderContext): @wp.kernel(module="unique", enable_backward=False) def _render_megakernel( # Model: - geom_type: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_matid: wp.array2d(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array2d(dtype=wp.vec4), - cam_projection: wp.array(dtype=int), - cam_fovy: wp.array2d(dtype=float), - cam_sensorsize: wp.array(dtype=wp.vec2), - cam_intrinsic: wp.array2d(dtype=wp.vec4), - light_type: wp.array2d(dtype=int), - light_castshadow: wp.array2d(dtype=bool), - light_active: wp.array2d(dtype=bool), - flex_vertadr: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_radius: wp.array(dtype=float), - mesh_faceadr: wp.array(dtype=int), - mat_texid: wp.array3d(dtype=int), - mat_texrepeat: wp.array2d(dtype=wp.vec2), - mat_rgba: wp.array2d(dtype=wp.vec4), + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_matid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + geom_rgba: wp.array2d[wp.vec4], + cam_projection: wp.array[int], + cam_fovy: wp.array2d[float], + cam_sensorsize: wp.array[wp.vec2], + cam_intrinsic: wp.array2d[wp.vec4], + light_type: wp.array2d[int], + light_castshadow: wp.array2d[bool], + light_active: wp.array2d[bool], + flex_vertadr: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], + mesh_faceadr: wp.array[int], + mat_texid: wp.array3d[int], + mat_texrepeat: wp.array2d[wp.vec2], + mat_rgba: wp.array2d[wp.vec4], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), - light_xpos_in: wp.array2d(dtype=wp.vec3), - light_xdir_in: wp.array2d(dtype=wp.vec3), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], + light_xpos_in: wp.array2d[wp.vec3], + light_xdir_in: wp.array2d[wp.vec3], + flexvert_xpos_in: wp.array2d[wp.vec3], # In: nrender: int, use_shadows: bool, bvh_ngeom: int, bvh_nflexgeom: int, - cam_res: wp.array(dtype=wp.vec2i), - cam_id_map: wp.array(dtype=int), - ray: wp.array(dtype=wp.vec3), - rgb_adr: wp.array(dtype=int), - depth_adr: wp.array(dtype=int), - seg_adr: wp.array(dtype=int), - render_rgb: wp.array(dtype=bool), - render_depth: wp.array(dtype=bool), - render_seg: wp.array(dtype=bool), + cam_res: wp.array[wp.vec2i], + cam_id_map: wp.array[int], + ray: wp.array[wp.vec3], + rgb_adr: wp.array[int], + depth_adr: wp.array[int], + seg_adr: wp.array[int], + render_rgb: wp.array[bool], + render_depth: wp.array[bool], + render_seg: wp.array[bool], bvh_id: wp.uint64, - group_root: wp.array(dtype=int), - flex_bvh_id: wp.array(dtype=wp.uint64), - flex_group_root: wp.array2d(dtype=int), - enabled_geom_ids: wp.array(dtype=int), - mesh_bvh_id: wp.array(dtype=wp.uint64), - mesh_facetexcoord: wp.array(dtype=wp.vec3i), - mesh_texcoord: wp.array(dtype=wp.vec2), - mesh_texcoord_offsets: wp.array(dtype=int), - hfield_bvh_id: wp.array(dtype=wp.uint64), - flex_rgba: wp.array(dtype=wp.vec4), - flex_geom_flexid: wp.array(dtype=int), - flex_geom_edgeid: wp.array(dtype=int), - textures: wp.array(dtype=wp.Texture2D), + group_root: wp.array[int], + flex_bvh_id: wp.array[wp.uint64], + flex_group_root: wp.array2d[int], + enabled_geom_ids: wp.array[int], + mesh_bvh_id: wp.array[wp.uint64], + mesh_facetexcoord: wp.array[wp.vec3i], + mesh_texcoord: wp.array[wp.vec2], + mesh_texcoord_offsets: wp.array[int], + hfield_bvh_id: wp.array[wp.uint64], + flex_rgba: wp.array[wp.vec4], + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], + textures: wp.array[wp.Texture2D], # Out: - rgb_out: wp.array2d(dtype=wp.uint32), - depth_out: wp.array2d(dtype=float), - seg_out: wp.array2d(dtype=int), + rgb_out: wp.array2d[wp.uint32], + depth_out: wp.array2d[float], + seg_out: wp.array2d[int], ): worldid, rayid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py index 36958f8e..c59da005 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py @@ -28,9 +28,9 @@ def _convert_texture_data( width: int, adr: int, nc: int, - tex_data_in: wp.array(dtype=wp.uint8), + tex_data_in: wp.array[wp.uint8], # Out: - tex_data_out: wp.array3d(dtype=float), + tex_data_out: wp.array3d[float], ): """Convert uint8 texture data to vec4 format for efficient sampling.""" 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 def unpack_rgb_kernel( # In: - packed: wp.array2d(dtype=wp.uint32), - rgb_adr: wp.array(dtype=int), + packed: wp.array2d[wp.uint32], + rgb_adr: wp.array[int], camera_index: int, # 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.""" worldid, pixelid = wp.tid() @@ -157,12 +157,12 @@ def unpack_rgb_kernel( @wp.kernel def extract_depth_kernel( # In: - depth_data: wp.array2d(dtype=float), - depth_adr: wp.array(dtype=int), + depth_data: wp.array2d[float], + depth_adr: wp.array[int], camera_index: int, depth_scale: float, # 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.""" 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) -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. 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. Args: @@ -211,11 +211,11 @@ def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_ou @wp.kernel def _extract_seg_kernel( # In: - seg_data: wp.array2d(dtype=int), - seg_adr: wp.array(dtype=int), + seg_data: wp.array2d[int], + seg_adr: wp.array[int], camera_index: int, # 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.""" worldid, pixelid = wp.tid() @@ -226,7 +226,7 @@ def _extract_seg_kernel( 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. Each pixel contains the MuJoCo geom ID of the geometry hit by the ray, -1 for diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py index 2c8177b8..b615d20c 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -53,15 +53,15 @@ wp.set_module_options({"enable_backward": False}) @wp.func def _write_scalar( # Model: - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], # In: sensorid: int, sensor: Any, # Out: - out: wp.array(dtype=float), + out: wp.array[float], ): adr = sensor_adr[sensorid] cutoff = sensor_cutoff[sensorid] @@ -81,16 +81,16 @@ def _write_scalar( @wp.func def _write_vector( # Model: - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], # In: sensorid: int, sensordim: int, sensor: Any, # Out: - out: wp.array(dtype=float), + out: wp.array[float], ): adr = sensor_adr[sensorid] cutoff = sensor_cutoff[sensorid] @@ -113,9 +113,9 @@ def _write_vector( @wp.func def _magnetometer( # Model: - opt_magnetic: wp.array(dtype=wp.vec3), + opt_magnetic: wp.array[wp.vec3], # Data in: - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, objid: int, @@ -127,14 +127,14 @@ def _magnetometer( @wp.func def _cam_projection( # Model: - cam_fovy: wp.array2d(dtype=float), - cam_resolution: wp.array(dtype=wp.vec2i), - cam_sensorsize: wp.array(dtype=wp.vec2), - cam_intrinsic: wp.array2d(dtype=wp.vec4), + cam_fovy: wp.array2d[float], + cam_resolution: wp.array[wp.vec2i], + cam_sensorsize: wp.array[wp.vec2], + cam_intrinsic: wp.array2d[wp.vec4], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), + site_xpos_in: wp.array2d[wp.vec3], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, objid: int, @@ -194,14 +194,14 @@ def _cam_projection( @wp.kernel def _sensor_rangefinder_init( # Model: - sensor_objid: wp.array(dtype=int), - sensor_rangefinder_adr: wp.array(dtype=int), + sensor_objid: wp.array[int], + sensor_rangefinder_adr: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], # Out: - pnt_out: wp.array2d(dtype=wp.vec3), - vec_out: wp.array2d(dtype=wp.vec3), + pnt_out: wp.array2d[wp.vec3], + vec_out: wp.array2d[wp.vec3], ): worldid, rfid = wp.tid() sensorid = sensor_rangefinder_adr[rfid] @@ -214,22 +214,22 @@ def _sensor_rangefinder_init( @wp.func -def _joint_pos(jnt_qposadr: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _joint_pos(jnt_qposadr: wp.array[int], qpos_in: wp.array2d[float], worldid: int, objid: int) -> float: return qpos_in[worldid, jnt_qposadr[objid]] @wp.func -def _tendon_pos(ten_length_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _tendon_pos(ten_length_in: wp.array2d[float], worldid: int, objid: int) -> float: return ten_length_in[worldid, objid] @wp.func -def _actuator_pos(actuator_length_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _actuator_pos(actuator_length_in: wp.array2d[float], worldid: int, objid: int) -> float: return actuator_length_in[worldid, objid] @wp.func -def _ball_quat(jnt_qposadr: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float), worldid: int, objid: int) -> wp.quat: +def _ball_quat(jnt_qposadr: wp.array[int], qpos_in: wp.array2d[float], worldid: int, objid: int) -> wp.quat: adr = jnt_qposadr[objid] quat = wp.quat( qpos_in[worldid, adr + 0], @@ -243,22 +243,22 @@ def _ball_quat(jnt_qposadr: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float @wp.kernel def _limit_pos( # Model: - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_limitpos_adr: wp.array(dtype=int), + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_objid: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_limitpos_adr: wp.array[int], # Data in: - ne_in: wp.array(dtype=int), - nf_in: wp.array(dtype=int), - nl_in: wp.array(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_pos_in: wp.array2d(dtype=float), - efc_margin_in: wp.array2d(dtype=float), + ne_in: wp.array[int], + nf_in: wp.array[int], + nl_in: wp.array[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_pos_in: wp.array2d[float], + efc_margin_in: wp.array2d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, efcid, limitposid = wp.tid() @@ -281,16 +281,16 @@ def _limit_pos( @wp.func def _frame_pos( # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, objid: int, @@ -340,11 +340,11 @@ def _frame_pos( @wp.func def _frame_axis( # Data in: - xmat_in: wp.array2d(dtype=wp.mat33), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xmat_in: wp.array2d(dtype=wp.mat33), + xmat_in: wp.array2d[wp.mat33], + ximat_in: wp.array2d[wp.mat33], + geom_xmat_in: wp.array2d[wp.mat33], + site_xmat_in: wp.array2d[wp.mat33], + cam_xmat_in: wp.array2d[wp.mat33], # In: worldid: int, objid: int, @@ -393,15 +393,15 @@ def _frame_axis( @wp.func def _frame_quat( # Model: - body_iquat: wp.array2d(dtype=wp.quat), - geom_bodyid: wp.array(dtype=int), - geom_quat: wp.array2d(dtype=wp.quat), - site_bodyid: wp.array(dtype=int), - site_quat: wp.array2d(dtype=wp.quat), - cam_bodyid: wp.array(dtype=int), - cam_quat: wp.array2d(dtype=wp.quat), + body_iquat: wp.array2d[wp.quat], + geom_bodyid: wp.array[int], + geom_quat: wp.array2d[wp.quat], + site_bodyid: wp.array[int], + site_quat: wp.array2d[wp.quat], + cam_bodyid: wp.array[int], + cam_quat: wp.array2d[wp.quat], # Data in: - xquat_in: wp.array2d(dtype=wp.quat), + xquat_in: wp.array2d[wp.quat], # In: worldid: int, objid: int, @@ -446,12 +446,12 @@ def _frame_quat( @wp.func -def _subtree_com(subtree_com_in: wp.array2d(dtype=wp.vec3), worldid: int, objid: int) -> wp.vec3: +def _subtree_com(subtree_com_in: wp.array2d[wp.vec3], worldid: int, objid: int) -> wp.vec3: return subtree_com_in[worldid, objid] @wp.func -def _clock(time_in: wp.array(dtype=float), worldid: int) -> float: +def _clock(time_in: wp.array[float], worldid: int) -> float: return time_in[worldid] @@ -459,58 +459,58 @@ def _clock(time_in: wp.array(dtype=float), worldid: int) -> float: def _sensor_pos( # Model: ngeom: int, - opt_magnetic: wp.array(dtype=wp.vec3), - body_geomnum: wp.array(dtype=int), - body_geomadr: wp.array(dtype=int), - body_iquat: wp.array2d(dtype=wp.quat), - jnt_qposadr: wp.array(dtype=int), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_quat: wp.array2d(dtype=wp.quat), - site_type: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_size: wp.array(dtype=wp.vec3), - site_quat: wp.array2d(dtype=wp.quat), - cam_bodyid: wp.array(dtype=int), - cam_quat: wp.array2d(dtype=wp.quat), - cam_fovy: wp.array2d(dtype=float), - cam_resolution: wp.array(dtype=wp.vec2i), - cam_sensorsize: wp.array(dtype=wp.vec2), - cam_intrinsic: wp.array2d(dtype=wp.vec4), - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_objtype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_reftype: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - nxn_pairid: wp.array(dtype=wp.vec2i), - sensor_pos_adr: wp.array(dtype=int), - rangefinder_sensor_adr: wp.array(dtype=int), + opt_magnetic: wp.array[wp.vec3], + body_geomnum: wp.array[int], + body_geomadr: wp.array[int], + body_iquat: wp.array2d[wp.quat], + jnt_qposadr: wp.array[int], + geom_type: wp.array[int], + geom_bodyid: wp.array[int], + geom_quat: wp.array2d[wp.quat], + site_type: wp.array[int], + site_bodyid: wp.array[int], + site_size: wp.array[wp.vec3], + site_quat: wp.array2d[wp.quat], + cam_bodyid: wp.array[int], + cam_quat: wp.array2d[wp.quat], + cam_fovy: wp.array2d[float], + cam_resolution: wp.array[wp.vec2i], + cam_sensorsize: wp.array[wp.vec2], + cam_intrinsic: wp.array2d[wp.vec4], + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_objtype: wp.array[int], + sensor_objid: wp.array[int], + sensor_reftype: wp.array[int], + sensor_refid: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + nxn_pairid: wp.array[wp.vec2i], + sensor_pos_adr: wp.array[int], + rangefinder_sensor_adr: wp.array[int], # Data in: - time_in: wp.array(dtype=float), - energy_in: wp.array(dtype=wp.vec2), - qpos_in: wp.array2d(dtype=float), - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), - xmat_in: wp.array2d(dtype=wp.mat33), - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - ten_length_in: wp.array2d(dtype=float), - actuator_length_in: wp.array2d(dtype=float), + time_in: wp.array[float], + energy_in: wp.array[wp.vec2], + qpos_in: wp.array2d[float], + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], + xmat_in: wp.array2d[wp.mat33], + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + ten_length_in: wp.array2d[float], + actuator_length_in: wp.array2d[float], # In: - rangefinder_dist_in: wp.array2d(dtype=float), - sensor_collision_in: wp.array4d(dtype=float), + rangefinder_dist_in: wp.array2d[float], + sensor_collision_in: wp.array4d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, posid = wp.tid() sensorid = sensor_pos_adr[posid] @@ -710,18 +710,18 @@ def _sensor_pos( def _sensor_collision( # Model: ngeom: int, - nxn_pairid: wp.array(dtype=wp.vec2i), + nxn_pairid: wp.array[wp.vec2i], # Data in: - contact_dist_in: wp.array(dtype=float), - contact_pos_in: wp.array(dtype=wp.vec3), - contact_frame_in: wp.array(dtype=wp.mat33), - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_worldid_in: wp.array(dtype=int), - contact_type_in: wp.array(dtype=int), - contact_geomcollisionid_in: wp.array(dtype=int), - nacon_in: wp.array(dtype=int), + contact_dist_in: wp.array[float], + contact_pos_in: wp.array[wp.vec3], + contact_frame_in: wp.array[wp.mat33], + contact_geom_in: wp.array[wp.vec2i], + contact_worldid_in: wp.array[int], + contact_type_in: wp.array[int], + contact_geomcollisionid_in: wp.array[int], + nacon_in: wp.array[int], # Out: - sensor_collision_out: wp.array4d(dtype=float), + sensor_collision_out: wp.array4d[float], ): conid = wp.tid() @@ -908,13 +908,13 @@ def sensor_pos(m: Model, d: Data): @wp.func def _velocimeter( # Model: - body_rootid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + site_bodyid: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -933,10 +933,10 @@ def _velocimeter( @wp.func def _gyro( # Model: - site_bodyid: wp.array(dtype=int), + site_bodyid: wp.array[int], # Data in: - site_xmat_in: wp.array2d(dtype=wp.mat33), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + site_xmat_in: wp.array2d[wp.mat33], + cvel_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -949,22 +949,22 @@ def _gyro( @wp.func -def _joint_vel(jnt_dofadr: wp.array(dtype=int), qvel_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _joint_vel(jnt_dofadr: wp.array[int], qvel_in: wp.array2d[float], worldid: int, objid: int) -> float: return qvel_in[worldid, jnt_dofadr[objid]] @wp.func -def _tendon_vel(ten_velocity_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _tendon_vel(ten_velocity_in: wp.array2d[float], worldid: int, objid: int) -> float: return ten_velocity_in[worldid, objid] @wp.func -def _actuator_vel(actuator_velocity_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _actuator_vel(actuator_velocity_in: wp.array2d[float], worldid: int, objid: int) -> float: return actuator_velocity_in[worldid, objid] @wp.func -def _ball_ang_vel(jnt_dofadr: wp.array(dtype=int), qvel_in: wp.array2d(dtype=float), worldid: int, objid: int) -> wp.vec3: +def _ball_ang_vel(jnt_dofadr: wp.array[int], qvel_in: wp.array2d[float], worldid: int, objid: int) -> wp.vec3: adr = jnt_dofadr[objid] return wp.vec3(qvel_in[worldid, adr + 0], qvel_in[worldid, adr + 1], qvel_in[worldid, adr + 2]) @@ -972,21 +972,21 @@ def _ball_ang_vel(jnt_dofadr: wp.array(dtype=int), qvel_in: wp.array2d(dtype=flo @wp.kernel def _limit_vel( # Model: - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_limitvel_adr: wp.array(dtype=int), + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_objid: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_limitvel_adr: wp.array[int], # Data in: - ne_in: wp.array(dtype=int), - nf_in: wp.array(dtype=int), - nl_in: wp.array(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_vel_in: wp.array2d(dtype=float), + ne_in: wp.array[int], + nf_in: wp.array[int], + nl_in: wp.array[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_vel_in: wp.array2d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, efcid, limitvelid = wp.tid() @@ -1010,18 +1010,18 @@ def _limit_vel( @wp.func def _cvel_offset( # Model: - body_rootid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xipos_in: wp.array2d(dtype=wp.vec3), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - site_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + xpos_in: wp.array2d[wp.vec3], + xipos_in: wp.array2d[wp.vec3], + geom_xpos_in: wp.array2d[wp.vec3], + site_xpos_in: wp.array2d[wp.vec3], + cam_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objtype: int, @@ -1052,23 +1052,23 @@ def _cvel_offset( @wp.func def _frame_linvel( # Model: - body_rootid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1158,23 +1158,23 @@ def _frame_linvel( @wp.func def _frame_angvel( # Model: - body_rootid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1238,52 +1238,52 @@ def _frame_angvel( @wp.func -def _subtree_linvel(subtree_linvel_in: wp.array2d(dtype=wp.vec3), worldid: int, objid: int) -> wp.vec3: +def _subtree_linvel(subtree_linvel_in: wp.array2d[wp.vec3], worldid: int, objid: int) -> wp.vec3: return subtree_linvel_in[worldid, objid] @wp.func -def _subtree_angmom(subtree_angmom_in: wp.array2d(dtype=wp.vec3), worldid: int, objid: int) -> wp.vec3: +def _subtree_angmom(subtree_angmom_in: wp.array2d[wp.vec3], worldid: int, objid: int) -> wp.vec3: return subtree_angmom_in[worldid, objid] @wp.kernel def _sensor_vel( # Model: - body_rootid: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_objtype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_reftype: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_vel_adr: wp.array(dtype=int), + body_rootid: wp.array[int], + jnt_dofadr: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_objtype: wp.array[int], + sensor_objid: wp.array[int], + sensor_reftype: wp.array[int], + sensor_refid: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_vel_adr: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - ten_velocity_in: wp.array2d(dtype=float), - actuator_velocity_in: wp.array2d(dtype=float), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - subtree_linvel_in: wp.array2d(dtype=wp.vec3), - subtree_angmom_in: wp.array2d(dtype=wp.vec3), + qvel_in: wp.array2d[float], + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + cam_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + ten_velocity_in: wp.array2d[float], + actuator_velocity_in: wp.array2d[float], + cvel_in: wp.array2d[wp.spatial_vector], + subtree_linvel_in: wp.array2d[wp.vec3], + subtree_angmom_in: wp.array2d[wp.vec3], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, velid = wp.tid() sensorid = sensor_vel_adr[velid] @@ -1450,14 +1450,14 @@ def sensor_vel(m: Model, d: Data): @wp.func def _accelerometer( # Model: - body_rootid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + site_bodyid: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - cacc_in: wp.array2d(dtype=wp.spatial_vector), + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], + cacc_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1482,10 +1482,10 @@ def _accelerometer( @wp.func def _force( # Model: - site_bodyid: wp.array(dtype=int), + site_bodyid: wp.array[int], # Data in: - site_xmat_in: wp.array2d(dtype=wp.mat33), - cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + site_xmat_in: wp.array2d[wp.mat33], + cfrc_int_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1499,13 +1499,13 @@ def _force( @wp.func def _torque( # Model: - body_rootid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + site_bodyid: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cfrc_int_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1518,16 +1518,16 @@ def _torque( @wp.func -def _actuator_force(actuator_force_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float: +def _actuator_force(actuator_force_in: wp.array2d[float], worldid: int, objid: int) -> float: return actuator_force_in[worldid, objid] @wp.func def _joint_actuator_force( # Model: - jnt_dofadr: wp.array(dtype=int), + jnt_dofadr: wp.array[int], # Data in: - qfrc_actuator_in: wp.array2d(dtype=float), + qfrc_actuator_in: wp.array2d[float], # In: worldid: int, objid: int, @@ -1538,15 +1538,15 @@ def _joint_actuator_force( @wp.kernel def _tendon_actuator_force( # Model: - actuator_trntype: wp.array(dtype=int), - actuator_trnid: wp.array(dtype=wp.vec2i), - sensor_objid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_tendonactfrc_adr: wp.array(dtype=int), + actuator_trntype: wp.array[int], + actuator_trnid: wp.array[wp.vec2i], + sensor_objid: wp.array[int], + sensor_adr: wp.array[int], + sensor_tendonactfrc_adr: wp.array[int], # Data in: - actuator_force_in: wp.array2d(dtype=float), + actuator_force_in: wp.array2d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, tenactfrcid, actid = wp.tid() sensorid = sensor_tendonactfrc_adr[tenactfrcid] @@ -1559,15 +1559,15 @@ def _tendon_actuator_force( @wp.kernel def _tendon_actuator_force_cutoff( # Model: - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_tendonactfrc_adr: wp.array(dtype=int), + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_tendonactfrc_adr: wp.array[int], # Data in: - sensordata_in: wp.array2d(dtype=float), + sensordata_in: wp.array2d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, tenactfrcid = wp.tid() sensorid = sensor_tendonactfrc_adr[tenactfrcid] @@ -1580,21 +1580,21 @@ def _tendon_actuator_force_cutoff( @wp.kernel def _limit_frc( # Model: - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_limitfrc_adr: wp.array(dtype=int), + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_objid: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_limitfrc_adr: wp.array[int], # Data in: - ne_in: wp.array(dtype=int), - nf_in: wp.array(dtype=int), - nl_in: wp.array(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_force_in: wp.array2d(dtype=float), + ne_in: wp.array[int], + nf_in: wp.array[int], + nl_in: wp.array[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_force_in: wp.array2d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, efcid, limitfrcid = wp.tid() @@ -1618,19 +1618,19 @@ def _limit_frc( @wp.func def _framelinacc( # Model: - body_rootid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xipos_in: wp.array2d(dtype=wp.vec3), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - site_xpos_in: wp.array2d(dtype=wp.vec3), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - cacc_in: wp.array2d(dtype=wp.spatial_vector), + xpos_in: wp.array2d[wp.vec3], + xipos_in: wp.array2d[wp.vec3], + geom_xpos_in: wp.array2d[wp.vec3], + site_xpos_in: wp.array2d[wp.vec3], + cam_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], + cacc_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1669,11 +1669,11 @@ def _framelinacc( @wp.func def _frameangacc( # Model: - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], # Data in: - cacc_in: wp.array2d(dtype=wp.spatial_vector), + cacc_in: wp.array2d[wp.spatial_vector], # In: worldid: int, objid: int, @@ -1697,49 +1697,49 @@ def _frameangacc( def _sensor_acc( # Model: opt_cone: int, - body_rootid: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), - sensor_type: wp.array(dtype=int), - sensor_datatype: wp.array(dtype=int), - sensor_objtype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_intprm: wp.array2d(dtype=int), - sensor_dim: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_acc_adr: wp.array(dtype=int), - sensor_adr_to_contact_adr: wp.array(dtype=int), + body_rootid: wp.array[int], + jnt_dofadr: wp.array[int], + geom_bodyid: wp.array[int], + site_bodyid: wp.array[int], + cam_bodyid: wp.array[int], + sensor_type: wp.array[int], + sensor_datatype: wp.array[int], + sensor_objtype: wp.array[int], + sensor_objid: wp.array[int], + sensor_intprm: wp.array2d[int], + sensor_dim: wp.array[int], + sensor_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_acc_adr: wp.array[int], + sensor_adr_to_contact_adr: wp.array[int], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xipos_in: wp.array2d(dtype=wp.vec3), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - cam_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - actuator_force_in: wp.array2d(dtype=float), - qfrc_actuator_in: wp.array2d(dtype=float), - cacc_in: wp.array2d(dtype=wp.spatial_vector), - cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), - contact_dist_in: wp.array(dtype=float), - contact_pos_in: wp.array(dtype=wp.vec3), - contact_frame_in: wp.array(dtype=wp.mat33), - contact_friction_in: wp.array(dtype=vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - efc_force_in: wp.array2d(dtype=float), + xpos_in: wp.array2d[wp.vec3], + xipos_in: wp.array2d[wp.vec3], + geom_xpos_in: wp.array2d[wp.vec3], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + cam_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], + actuator_force_in: wp.array2d[float], + qfrc_actuator_in: wp.array2d[float], + cacc_in: wp.array2d[wp.spatial_vector], + cfrc_int_in: wp.array2d[wp.spatial_vector], + contact_dist_in: wp.array[float], + contact_pos_in: wp.array[wp.vec3], + contact_frame_in: wp.array[wp.mat33], + contact_friction_in: wp.array[vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + efc_force_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - sensor_contact_nmatch_in: wp.array2d(dtype=int), - sensor_contact_matchid_in: wp.array3d(dtype=int), - sensor_contact_direction_in: wp.array3d(dtype=float), + sensor_contact_nmatch_in: wp.array2d[int], + sensor_contact_matchid_in: wp.array3d[int], + sensor_contact_direction_in: wp.array3d[float], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, accid = wp.tid() sensorid = sensor_acc_adr[accid] @@ -2001,26 +2001,26 @@ def _sensor_acc( def _sensor_touch( # Model: opt_cone: int, - geom_bodyid: wp.array(dtype=int), - site_type: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_size: wp.array(dtype=wp.vec3), - sensor_objid: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_touch_adr: wp.array(dtype=int), + geom_bodyid: wp.array[int], + site_type: wp.array[int], + site_bodyid: wp.array[int], + site_size: wp.array[wp.vec3], + sensor_objid: wp.array[int], + sensor_adr: wp.array[int], + sensor_touch_adr: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - contact_pos_in: wp.array(dtype=wp.vec3), - contact_frame_in: wp.array(dtype=wp.mat33), - contact_dim_in: wp.array(dtype=int), - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - efc_force_in: wp.array2d(dtype=float), - nacon_in: wp.array(dtype=int), + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + contact_pos_in: wp.array[wp.vec3], + contact_frame_in: wp.array[wp.mat33], + contact_dim_in: wp.array[int], + contact_geom_in: wp.array[wp.vec2i], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + efc_force_in: wp.array2d[float], + nacon_in: wp.array[int], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): conid, sensortouchadrid = wp.tid() @@ -2084,15 +2084,15 @@ def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3: @wp.kernel def _preprocess_tactile_contacts( # Model: - body_weldid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), + body_weldid: wp.array[int], + geom_bodyid: wp.array[int], # Data in: - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_worldid_in: wp.array(dtype=int), - nacon_in: wp.array(dtype=int), + contact_geom_in: wp.array[wp.vec2i], + contact_worldid_in: wp.array[int], + nacon_in: wp.array[int], # Out: - weld_geom_count_out: wp.array2d(dtype=int), - weld_geom_list_out: wp.array3d(dtype=int), + weld_geom_count_out: wp.array2d[int], + weld_geom_list_out: wp.array3d[int], ): conid = wp.tid() ncon = nacon_in[0] @@ -2121,42 +2121,42 @@ def _preprocess_tactile_contacts( @wp.kernel def _sensor_tactile( # Model: - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - oct_child: wp.array(dtype=vec8i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_octadr: wp.array(dtype=int), - mesh_normaladr: wp.array(dtype=int), - mesh_normalnum: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_normal: wp.array(dtype=wp.vec3), - mesh_quat: wp.array(dtype=wp.quat), - sensor_objid: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_dim: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=vec_pluginattr), - geom_plugin_index: wp.array(dtype=int), - taxel_vertadr: wp.array(dtype=int), - taxel_sensorid: wp.array(dtype=int), + body_rootid: wp.array[int], + body_weldid: wp.array[int], + oct_child: wp.array[vec8i], + oct_aabb: wp.array2d[wp.vec3], + oct_coeff: wp.array[vec8], + geom_type: wp.array[int], + geom_bodyid: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], + mesh_octadr: wp.array[int], + mesh_normaladr: wp.array[int], + mesh_normalnum: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_normal: wp.array[wp.vec3], + mesh_quat: wp.array[wp.quat], + sensor_objid: wp.array[int], + sensor_refid: wp.array[int], + sensor_dim: wp.array[int], + sensor_adr: wp.array[int], + plugin: wp.array[int], + plugin_attr: wp.array[vec_pluginattr], + geom_plugin_index: wp.array[int], + taxel_vertadr: wp.array[int], + taxel_sensorid: wp.array[int], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # In: - weld_geom_count_in: wp.array2d(dtype=int), - weld_geom_list_in: wp.array3d(dtype=int), + weld_geom_count_in: wp.array2d[int], + weld_geom_list_in: wp.array3d[int], # Data out: - sensordata_out: wp.array2d(dtype=float), + sensordata_out: wp.array2d[float], ): worldid, taxelid = wp.tid() @@ -2223,7 +2223,7 @@ def _sensor_tactile( contact_type, geom_size[worldid % geom_size.shape[0], geom], plugin_id, - geom_dataid[geom], + geom_dataid[worldid % geom_dataid.shape[0], geom], ) depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0) @@ -2253,7 +2253,7 @@ def _sensor_tactile( @wp.func -def _check_match(body_parentid: wp.array(dtype=int), body: int, geom: int, objtype: int, objid: int) -> bool: +def _check_match(body_parentid: wp.array[int], body: int, geom: int, objtype: int, objid: int) -> bool: """Check if a contact body/geom matches a sensor spec (objtype, objid).""" if objtype == ObjType.UNKNOWN: return True @@ -2276,36 +2276,36 @@ def _contact_match( # Model: opt_cone: int, opt_contact_sensor_maxmatch: int, - body_parentid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - site_type: wp.array(dtype=int), - site_size: wp.array(dtype=wp.vec3), - sensor_objtype: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_reftype: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_intprm: wp.array2d(dtype=int), - sensor_contact_adr: wp.array(dtype=int), + body_parentid: wp.array[int], + geom_bodyid: wp.array[int], + site_type: wp.array[int], + site_size: wp.array[wp.vec3], + sensor_objtype: wp.array[int], + sensor_objid: wp.array[int], + sensor_reftype: wp.array[int], + sensor_refid: wp.array[int], + sensor_intprm: wp.array2d[int], + sensor_contact_adr: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - contact_dist_in: wp.array(dtype=float), - contact_pos_in: wp.array(dtype=wp.vec3), - contact_frame_in: wp.array(dtype=wp.mat33), - contact_friction_in: wp.array(dtype=vec5), - contact_dim_in: wp.array(dtype=int), - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - contact_type_in: wp.array(dtype=int), - efc_force_in: wp.array2d(dtype=float), + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + contact_dist_in: wp.array[float], + contact_pos_in: wp.array[wp.vec3], + contact_frame_in: wp.array[wp.mat33], + contact_friction_in: wp.array[vec5], + contact_dim_in: wp.array[int], + contact_geom_in: wp.array[wp.vec2i], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + contact_type_in: wp.array[int], + efc_force_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # Out: - sensor_contact_nmatch_out: wp.array2d(dtype=int), - sensor_contact_matchid_out: wp.array3d(dtype=int), - sensor_contact_criteria_out: wp.array3d(dtype=float), - sensor_contact_direction_out: wp.array3d(dtype=float), + sensor_contact_nmatch_out: wp.array2d[int], + sensor_contact_matchid_out: wp.array3d[int], + sensor_contact_criteria_out: wp.array3d[float], + sensor_contact_direction_out: wp.array3d[float], ): contactsensorid, contactid = wp.tid() sensorid = sensor_contact_adr[contactsensorid] @@ -2411,14 +2411,14 @@ def _contact_sort(maxmatch: int): @wp.kernel(module="unique", enable_backward=False) def contact_sort( # Model: - sensor_intprm: wp.array2d(dtype=int), - sensor_contact_adr: wp.array(dtype=int), + sensor_intprm: wp.array2d[int], + sensor_contact_adr: wp.array[int], # In: - sensor_contact_nmatch_in: wp.array2d(dtype=int), - sensor_contact_matchid_in: wp.array3d(dtype=int), - sensor_contact_criteria_in: wp.array3d(dtype=float), + sensor_contact_nmatch_in: wp.array2d[int], + sensor_contact_matchid_in: wp.array3d[int], + sensor_contact_criteria_in: wp.array3d[float], # Out: - sensor_contact_matchid_out: wp.array3d(dtype=int), + sensor_contact_matchid_out: wp.array3d[int], ): worldid, contactsensorid = wp.tid() @@ -2700,7 +2700,7 @@ def sensor_acc(m: Model, d: Data): @wp.kernel def _energy_pos_zero( # Data out: - energy_out: wp.array(dtype=wp.vec2), + energy_out: wp.array[wp.vec2], ): worldid = wp.tid() energy_out[worldid][0] = 0.0 @@ -2709,12 +2709,12 @@ def _energy_pos_zero( @wp.kernel def _energy_pos_gravity( # Model: - opt_gravity: wp.array(dtype=wp.vec3), - body_mass: wp.array2d(dtype=float), + opt_gravity: wp.array[wp.vec3], + body_mass: wp.array2d[float], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), + xipos_in: wp.array2d[wp.vec3], # Data out: - energy_out: wp.array(dtype=wp.vec2), + energy_out: wp.array[wp.vec2], ): worldid, bodyid = wp.tid() gravity = opt_gravity[worldid % opt_gravity.shape[0]] @@ -2731,14 +2731,14 @@ def _energy_pos_gravity( @wp.kernel def _energy_pos_passive_joint( # Model: - qpos_spring: wp.array2d(dtype=float), - jnt_type: wp.array(dtype=int), - jnt_qposadr: wp.array(dtype=int), - jnt_stiffness: wp.array2d(dtype=float), + qpos_spring: wp.array2d[float], + jnt_type: wp.array[int], + jnt_qposadr: wp.array[int], + jnt_stiffness: wp.array2d[float], # Data in: - qpos_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], # Data out: - energy_out: wp.array(dtype=wp.vec2), + energy_out: wp.array[wp.vec2], ): worldid, jntid = wp.tid() jnt_stiffness_id = worldid % jnt_stiffness.shape[0] @@ -2817,12 +2817,12 @@ def _energy_pos_passive_joint( @wp.kernel def _energy_pos_passive_tendon( # Model: - tendon_stiffness: wp.array2d(dtype=float), - tendon_lengthspring: wp.array2d(dtype=wp.vec2), + tendon_stiffness: wp.array2d[float], + tendon_lengthspring: wp.array2d[wp.vec2], # Data in: - ten_length_in: wp.array2d(dtype=float), + ten_length_in: wp.array2d[float], # Data out: - energy_out: wp.array(dtype=wp.vec2), + energy_out: wp.array[wp.vec2], ): worldid, tenid = wp.tid() @@ -2897,11 +2897,11 @@ def _energy_vel_kinetic(nv: int): @wp.kernel(module="unique", enable_backward=False) def energy_vel_kinetic( # Data in: - qvel_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], # In: - Mqvel: wp.array2d(dtype=float), + Mqvel: wp.array2d[float], # Data out: - energy_out: wp.array(dtype=wp.vec2), + energy_out: wp.array[wp.vec2], ): worldid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py index 51dfadb4..bdb90ecd 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py @@ -44,28 +44,28 @@ wp.set_module_options({"enable_backward": False}) @wp.kernel def _kinematics_branch( # Model: - qpos0: wp.array2d(dtype=float), - body_parentid: wp.array(dtype=int), - body_mocapid: wp.array(dtype=int), - body_jntnum: wp.array(dtype=int), - body_jntadr: wp.array(dtype=int), - body_pos: wp.array2d(dtype=wp.vec3), - body_quat: wp.array2d(dtype=wp.quat), - jnt_type: wp.array(dtype=int), - jnt_qposadr: wp.array(dtype=int), - jnt_pos: wp.array2d(dtype=wp.vec3), - jnt_axis: wp.array2d(dtype=wp.vec3), - body_branches: wp.array(dtype=int), - body_branch_start: wp.array(dtype=int), + qpos0: wp.array2d[float], + body_parentid: wp.array[int], + body_mocapid: wp.array[int], + body_jntnum: wp.array[int], + body_jntadr: wp.array[int], + body_pos: wp.array2d[wp.vec3], + body_quat: wp.array2d[wp.quat], + jnt_type: wp.array[int], + jnt_qposadr: wp.array[int], + jnt_pos: wp.array2d[wp.vec3], + jnt_axis: wp.array2d[wp.vec3], + body_branches: wp.array[int], + body_branch_start: wp.array[int], # Data in: - qpos_in: wp.array2d(dtype=float), - mocap_pos_in: wp.array2d(dtype=wp.vec3), - mocap_quat_in: wp.array2d(dtype=wp.quat), + qpos_in: wp.array2d[float], + mocap_pos_in: wp.array2d[wp.vec3], + mocap_quat_in: wp.array2d[wp.quat], # Data out: - xpos_out: wp.array2d(dtype=wp.vec3), - xquat_out: wp.array2d(dtype=wp.quat), - xanchor_out: wp.array2d(dtype=wp.vec3), - xaxis_out: wp.array2d(dtype=wp.vec3), + xpos_out: wp.array2d[wp.vec3], + xquat_out: wp.array2d[wp.quat], + xanchor_out: wp.array2d[wp.vec3], + xaxis_out: wp.array2d[wp.vec3], ): worldid, branchid = wp.tid() @@ -146,14 +146,14 @@ def _kinematics_branch( @wp.kernel def _compute_body_inertial_frames( # Model: - body_ipos: wp.array2d(dtype=wp.vec3), - body_iquat: wp.array2d(dtype=wp.quat), + body_ipos: wp.array2d[wp.vec3], + body_iquat: wp.array2d[wp.quat], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], # Data out: - xipos_out: wp.array2d(dtype=wp.vec3), - ximat_out: wp.array2d(dtype=wp.mat33), + xipos_out: wp.array2d[wp.vec3], + ximat_out: wp.array2d[wp.mat33], ): worldid, bodyid = wp.tid() xpos = xpos_in[worldid, bodyid] @@ -165,9 +165,9 @@ def _compute_body_inertial_frames( @wp.kernel def _compute_body_matrices( # Data in: - xquat_in: wp.array2d(dtype=wp.quat), + xquat_in: wp.array2d[wp.quat], # Data out: - xmat_out: wp.array2d(dtype=wp.mat33), + xmat_out: wp.array2d[wp.mat33], ): worldid, bodyid = wp.tid() xmat_out[worldid, bodyid] = math.quat_to_mat(xquat_in[worldid, bodyid]) @@ -176,18 +176,18 @@ def _compute_body_matrices( @wp.kernel def _geom_local_to_global( # Model: - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - body_mocapid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_pos: wp.array2d(dtype=wp.vec3), - geom_quat: wp.array2d(dtype=wp.quat), + body_rootid: wp.array[int], + body_weldid: wp.array[int], + body_mocapid: wp.array[int], + geom_bodyid: wp.array[int], + geom_pos: wp.array2d[wp.vec3], + geom_quat: wp.array2d[wp.quat], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], # Data out: - geom_xpos_out: wp.array2d(dtype=wp.vec3), - geom_xmat_out: wp.array2d(dtype=wp.mat33), + geom_xpos_out: wp.array2d[wp.vec3], + geom_xmat_out: wp.array2d[wp.mat33], ): worldid, geomid = wp.tid() bodyid = geom_bodyid[geomid] @@ -206,15 +206,15 @@ def _geom_local_to_global( @wp.kernel def _site_local_to_global( # Model: - site_bodyid: wp.array(dtype=int), - site_pos: wp.array2d(dtype=wp.vec3), - site_quat: wp.array2d(dtype=wp.quat), + site_bodyid: wp.array[int], + site_pos: wp.array2d[wp.vec3], + site_quat: wp.array2d[wp.quat], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], # Data out: - site_xpos_out: wp.array2d(dtype=wp.vec3), - site_xmat_out: wp.array2d(dtype=wp.mat33), + site_xpos_out: wp.array2d[wp.vec3], + site_xmat_out: wp.array2d[wp.mat33], ): worldid, siteid = wp.tid() bodyid = site_bodyid[siteid] @@ -228,16 +228,16 @@ def _site_local_to_global( def _flex_vertices( # Model: nflex: int, - flex_vertadr: wp.array(dtype=int), - flex_vertnum: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), - flex_vert: wp.array(dtype=wp.vec3), - flex_centered: wp.array(dtype=bool), + flex_vertadr: wp.array[int], + flex_vertnum: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_vert: wp.array[wp.vec3], + flex_centered: wp.array[bool], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], # Data out: - flexvert_xpos_out: wp.array2d(dtype=wp.vec3), + flexvert_xpos_out: wp.array2d[wp.vec3], ): worldid, vertid = wp.tid() @@ -261,25 +261,25 @@ def _flex_vertices( def _flex_edges( # Model: nflex: int, - body_rootid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_edgeadr: wp.array(dtype=int), - flex_edgenum: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flexedge_J_rowadr: wp.array(dtype=int), - flexedge_J_colind: wp.array(dtype=int), + body_rootid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + flex_vertadr: wp.array[int], + flex_edgeadr: wp.array[int], + flex_edgenum: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flexedge_J_rowadr: wp.array[int], + flexedge_J_colind: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - flexvert_xpos_in: wp.array2d(dtype=wp.vec3), + qvel_in: wp.array2d[float], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + flexvert_xpos_in: wp.array2d[wp.vec3], # Data out: - flexedge_J_out: wp.array2d(dtype=float), - flexedge_length_out: wp.array2d(dtype=float), - flexedge_velocity_out: wp.array2d(dtype=float), + flexedge_J_out: wp.array2d[float], + flexedge_length_out: wp.array2d[float], + flexedge_velocity_out: wp.array2d[float], ): worldid, edgeid = wp.tid() for i in range(nflex): @@ -463,11 +463,11 @@ def flex(m: Model, d: Data): @wp.kernel def _subtree_com_init( # Model: - body_mass: wp.array2d(dtype=float), + body_mass: wp.array2d[float], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), + xipos_in: wp.array2d[wp.vec3], # Data out: - subtree_com_out: wp.array2d(dtype=wp.vec3), + subtree_com_out: wp.array2d[wp.vec3], ): worldid, bodyid = wp.tid() subtree_com_out[worldid, bodyid] = xipos_in[worldid, bodyid] * body_mass[worldid % body_mass.shape[0], bodyid] @@ -476,13 +476,13 @@ def _subtree_com_init( @wp.kernel def _subtree_com_acc( # Model: - body_parentid: wp.array(dtype=int), + body_parentid: wp.array[int], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), + subtree_com_in: wp.array2d[wp.vec3], # In: - body_tree_: wp.array(dtype=int), + body_tree_: wp.array[int], # Data out: - subtree_com_out: wp.array2d(dtype=wp.vec3), + subtree_com_out: wp.array2d[wp.vec3], ): worldid, nodeid = wp.tid() bodyid = body_tree_[nodeid] @@ -494,11 +494,11 @@ def _subtree_com_acc( @wp.kernel def _subtree_div( # Model: - body_subtreemass: wp.array2d(dtype=float), + body_subtreemass: wp.array2d[float], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), + subtree_com_in: wp.array2d[wp.vec3], # Data out: - subtree_com_out: wp.array2d(dtype=wp.vec3), + subtree_com_out: wp.array2d[wp.vec3], ): worldid, bodyid = wp.tid() com = subtree_com_in[worldid, bodyid] @@ -510,15 +510,15 @@ def _subtree_div( @wp.kernel def _cinert( # Model: - body_rootid: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_inertia: wp.array2d(dtype=wp.vec3), + body_rootid: wp.array[int], + body_mass: wp.array2d[float], + body_inertia: wp.array2d[wp.vec3], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], # Data out: - cinert_out: wp.array2d(dtype=vec10), + cinert_out: wp.array2d[vec10], ): worldid, bodyid = wp.tid() mat = ximat_in[worldid, bodyid] @@ -556,17 +556,17 @@ def _cinert( @wp.kernel def _cdof( # Model: - body_rootid: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + jnt_type: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_bodyid: wp.array[int], # Data in: - xmat_in: wp.array2d(dtype=wp.mat33), - xanchor_in: wp.array2d(dtype=wp.vec3), - xaxis_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), + xmat_in: wp.array2d[wp.mat33], + xanchor_in: wp.array2d[wp.vec3], + xaxis_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], # Data out: - cdof_out: wp.array2d(dtype=wp.spatial_vector), + cdof_out: wp.array2d[wp.spatial_vector], ): worldid, jntid = wp.tid() bodyid = jnt_bodyid[jntid] @@ -635,21 +635,21 @@ def com_pos(m: Model, d: Data): @wp.kernel def _cam_local_to_global( # Model: - cam_mode: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), - cam_targetbodyid: wp.array(dtype=int), - cam_pos: wp.array2d(dtype=wp.vec3), - cam_quat: wp.array2d(dtype=wp.quat), - cam_poscom0: wp.array2d(dtype=wp.vec3), - cam_pos0: wp.array2d(dtype=wp.vec3), - cam_mat0: wp.array2d(dtype=wp.mat33), + cam_mode: wp.array[int], + cam_bodyid: wp.array[int], + cam_targetbodyid: wp.array[int], + cam_pos: wp.array2d[wp.vec3], + cam_quat: wp.array2d[wp.quat], + cam_poscom0: wp.array2d[wp.vec3], + cam_pos0: wp.array2d[wp.vec3], + cam_mat0: wp.array2d[wp.mat33], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), - subtree_com_in: wp.array2d(dtype=wp.vec3), + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], + subtree_com_in: wp.array2d[wp.vec3], # Data out: - cam_xpos_out: wp.array2d(dtype=wp.vec3), - cam_xmat_out: wp.array2d(dtype=wp.mat33), + cam_xpos_out: wp.array2d[wp.vec3], + cam_xmat_out: wp.array2d[wp.mat33], ): worldid, camid = wp.tid() cam_pos_id = worldid % cam_pos.shape[0] @@ -702,21 +702,21 @@ def _cam_local_to_global( @wp.kernel def _light_local_to_global( # Model: - light_mode: wp.array(dtype=int), - light_bodyid: wp.array(dtype=int), - light_targetbodyid: wp.array(dtype=int), - light_pos: wp.array2d(dtype=wp.vec3), - light_dir: wp.array2d(dtype=wp.vec3), - light_poscom0: wp.array2d(dtype=wp.vec3), - light_pos0: wp.array2d(dtype=wp.vec3), - light_dir0: wp.array2d(dtype=wp.vec3), + light_mode: wp.array[int], + light_bodyid: wp.array[int], + light_targetbodyid: wp.array[int], + light_pos: wp.array2d[wp.vec3], + light_dir: wp.array2d[wp.vec3], + light_poscom0: wp.array2d[wp.vec3], + light_pos0: wp.array2d[wp.vec3], + light_dir0: wp.array2d[wp.vec3], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), - subtree_com_in: wp.array2d(dtype=wp.vec3), + xpos_in: wp.array2d[wp.vec3], + xquat_in: wp.array2d[wp.quat], + subtree_com_in: wp.array2d[wp.vec3], # Data out: - light_xpos_out: wp.array2d(dtype=wp.vec3), - light_xdir_out: wp.array2d(dtype=wp.vec3), + light_xpos_out: wp.array2d[wp.vec3], + light_xdir_out: wp.array2d[wp.vec3], ): worldid, lightid = wp.tid() light_pos_id = worldid % light_pos.shape[0] @@ -806,13 +806,13 @@ def camlight(m: Model, d: Data): @wp.kernel def _crb_accumulate( # Model: - body_parentid: wp.array(dtype=int), + body_parentid: wp.array[int], # Data in: - crb_in: wp.array2d(dtype=vec10), + crb_in: wp.array2d[vec10], # In: - body_tree_: wp.array(dtype=int), + body_tree_: wp.array[int], # Data out: - crb_out: wp.array2d(dtype=vec10), + crb_out: wp.array2d[vec10], ): worldid, nodeid = wp.tid() bodyid = body_tree_[nodeid] @@ -825,15 +825,15 @@ def _crb_accumulate( @wp.kernel def _qM_sparse( # Model: - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - dof_Madr: wp.array(dtype=int), - dof_armature: wp.array2d(dtype=float), + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + dof_Madr: wp.array[int], + dof_armature: wp.array2d[float], # Data in: - cdof_in: wp.array2d(dtype=wp.spatial_vector), - crb_in: wp.array2d(dtype=vec10), + cdof_in: wp.array2d[wp.spatial_vector], + crb_in: wp.array2d[vec10], # Data out: - qM_out: wp.array3d(dtype=float), + qM_out: wp.array3d[float], ): worldid, dofid = wp.tid() madr_ij = dof_Madr[dofid] # dof_Madr is not batched @@ -855,14 +855,14 @@ def _qM_sparse( @wp.kernel def _qM_dense( # Model: - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - dof_armature: wp.array2d(dtype=float), + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + dof_armature: wp.array2d[float], # Data in: - cdof_in: wp.array2d(dtype=wp.spatial_vector), - crb_in: wp.array2d(dtype=vec10), + cdof_in: wp.array2d[wp.spatial_vector], + crb_in: wp.array2d[vec10], # Data out: - qM_out: wp.array3d(dtype=float), + qM_out: wp.array3d[float], ): worldid, dofid = wp.tid() bodyid = dof_bodyid[dofid] @@ -915,17 +915,17 @@ def crb(m: Model, d: Data): @wp.kernel def _tendon_armature( # Model: - dof_parentid: wp.array(dtype=int), - dof_Madr: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_armature: wp.array2d(dtype=float), + dof_parentid: wp.array[int], + dof_Madr: wp.array[int], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_armature: wp.array2d[float], is_sparse: bool, # Data in: - ten_J_in: wp.array2d(dtype=float), + ten_J_in: wp.array2d[float], # Data out: - qM_out: wp.array3d(dtype=float), + qM_out: wp.array3d[float], ): worldid, tenid, dofid = wp.tid() @@ -1003,11 +1003,11 @@ def tendon_armature(m: Model, d: Data): @wp.kernel def _copy_CSR( # Model: - mapM2M: wp.array(dtype=int), + mapM2M: wp.array[int], # In: - M_in: wp.array3d(dtype=float), + M_in: wp.array3d[float], # Out: - L_out: wp.array3d(dtype=float), + L_out: wp.array3d[float], ): worldid, ind = wp.tid() L_out[worldid, 0, ind] = M_in[worldid, 0, mapM2M[ind]] @@ -1016,13 +1016,13 @@ def _copy_CSR( @wp.kernel def _qLD_acc( # Model: - M_rownnz: wp.array(dtype=int), - M_rowadr: wp.array(dtype=int), + M_rownnz: wp.array[int], + M_rowadr: wp.array[int], # In: - qLD_updates_: wp.array(dtype=wp.vec3i), - L_in: wp.array3d(dtype=float), + qLD_updates_: wp.array[wp.vec3i], + L_in: wp.array3d[float], # Out: - L_out: wp.array3d(dtype=float), + L_out: wp.array3d[float], ): worldid, nodeid = wp.tid() update = qLD_updates_[nodeid] @@ -1041,19 +1041,19 @@ def _qLD_acc( @wp.kernel def _qLDiag_div( # Model: - M_rownnz: wp.array(dtype=int), - M_rowadr: wp.array(dtype=int), + M_rownnz: wp.array[int], + M_rowadr: wp.array[int], # In: - L_in: wp.array3d(dtype=float), + L_in: wp.array3d[float], # Out: - D_out: wp.array2d(dtype=float), + D_out: wp.array2d[float], ): worldid, dofid = wp.tid() diag_i = M_rowadr[dofid] + M_rownnz[dofid] - 1 # Address of diagonal element of i D_out[worldid, dofid] = 1.0 / L_in[worldid, 0, diag_i] -def _factor_i_sparse(m: Model, d: Data, M: wp.array3d(dtype=float), L: wp.array3d(dtype=float), D: wp.array2d(dtype=float)): +def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[float], D: wp.array2d[float]): """Sparse L'*D*L factorization of inertia-like matrix M, assumed spd.""" wp.launch(_copy_CSR, dim=(d.nworld, m.nC), inputs=[m.mapM2M, M], outputs=[L]) @@ -1071,11 +1071,11 @@ def _tile_cholesky_factorize(tile: TileSet): @wp.kernel(module="unique", enable_backward=False) def cholesky_factorize( # Data in: - qM_in: wp.array3d(dtype=float), + qM_in: wp.array3d[float], # In: - adr: wp.array(dtype=int), + adr: wp.array[int], # Out: - L_out: wp.array3d(dtype=float), + L_out: wp.array3d[float], ): worldid, nodeid = wp.tid() TILE_SIZE = wp.static(tile.size) @@ -1112,9 +1112,9 @@ def factor_m(m: Model, d: Data): @wp.kernel def _cacc_world( # In: - gravity: wp.array(dtype=wp.vec3), + gravity: wp.array[wp.vec3], # Data out: - cacc_out: wp.array2d(dtype=wp.spatial_vector), + cacc_out: wp.array2d[wp.spatial_vector], ): worldid = wp.tid() cacc_out[worldid, 0] = wp.spatial_vector(wp.vec3(0.0), -gravity[worldid % gravity.shape[0]]) @@ -1130,20 +1130,20 @@ def _rne_cacc_world(m: Model, d: Data): @wp.kernel def _cacc_branch( # Model: - body_parentid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_branches: wp.array(dtype=int), - body_branch_start: wp.array(dtype=int), + body_parentid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + body_branches: wp.array[int], + body_branch_start: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - qacc_in: wp.array2d(dtype=float), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - cdof_dot_in: wp.array2d(dtype=wp.spatial_vector), + qvel_in: wp.array2d[float], + qacc_in: wp.array2d[float], + cdof_in: wp.array2d[wp.spatial_vector], + cdof_dot_in: wp.array2d[wp.spatial_vector], # In: flg_acc: bool, # Data out: - cacc_out: wp.array2d(dtype=wp.spatial_vector), + cacc_out: wp.array2d[wp.spatial_vector], ): worldid, branchid = wp.tid() @@ -1187,14 +1187,14 @@ def _rne_cacc_forward(m: Model, d: Data, flg_acc: bool = False): @wp.kernel def _cfrc( # Data in: - cinert_in: wp.array2d(dtype=vec10), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - cacc_in: wp.array2d(dtype=wp.spatial_vector), - cfrc_ext_in: wp.array2d(dtype=wp.spatial_vector), + cinert_in: wp.array2d[vec10], + cvel_in: wp.array2d[wp.spatial_vector], + cacc_in: wp.array2d[wp.spatial_vector], + cfrc_ext_in: wp.array2d[wp.spatial_vector], # In: flg_cfrc_ext: bool, # Data out: - cfrc_int_out: wp.array2d(dtype=wp.spatial_vector), + cfrc_int_out: wp.array2d[wp.spatial_vector], ): worldid, bodyid = wp.tid() bodyid += 1 # skip world body @@ -1218,13 +1218,13 @@ def _rne_cfrc(m: Model, d: Data, flg_cfrc_ext: bool = False): @wp.kernel def _cfrc_backward( # Model: - body_parentid: wp.array(dtype=int), + body_parentid: wp.array[int], # Data in: - cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + cfrc_int_in: wp.array2d[wp.spatial_vector], # In: - body_tree_: wp.array(dtype=int), + body_tree_: wp.array[int], # Data out: - cfrc_int_out: wp.array2d(dtype=wp.spatial_vector), + cfrc_int_out: wp.array2d[wp.spatial_vector], ): worldid, nodeid = wp.tid() bodyid = body_tree_[nodeid] @@ -1243,12 +1243,12 @@ def _rne_cfrc_backward(m: Model, d: Data): @wp.kernel def _qfrc_bias( # Model: - dof_bodyid: wp.array(dtype=int), + dof_bodyid: wp.array[int], # Data in: - cdof_in: wp.array2d(dtype=wp.spatial_vector), - cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + cdof_in: wp.array2d[wp.spatial_vector], + cfrc_int_in: wp.array2d[wp.spatial_vector], # Data out: - qfrc_bias_out: wp.array2d(dtype=float), + qfrc_bias_out: wp.array2d[float], ): worldid, dofid = wp.tid() bodyid = dof_bodyid[dofid] @@ -1277,13 +1277,13 @@ def rne(m: Model, d: Data, flg_acc: bool = False): @wp.kernel def _cfrc_ext( # Model: - body_rootid: wp.array(dtype=int), + body_rootid: wp.array[int], # Data in: - xfrc_applied_in: wp.array2d(dtype=wp.spatial_vector), - xipos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), + xfrc_applied_in: wp.array2d[wp.spatial_vector], + xipos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], # Data out: - cfrc_ext_out: wp.array2d(dtype=wp.spatial_vector), + cfrc_ext_out: wp.array2d[wp.spatial_vector], ): worldid, bodyid = wp.tid() if bodyid == 0: @@ -1298,14 +1298,14 @@ def _cfrc_ext( @wp.kernel def _count_equality_constraints( # Model: - eq_type: wp.array(dtype=int), + eq_type: wp.array[int], # Data in: - ne_in: wp.array(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), + ne_in: wp.array[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], # Out: - ne_connect_out: wp.array(dtype=int), - ne_weld_out: wp.array(dtype=int), + ne_connect_out: wp.array[int], + ne_weld_out: wp.array[int], ): """Counts connect and weld equality constraints from efc data.""" worldid, efcid = wp.tid() @@ -1328,24 +1328,24 @@ def _count_equality_constraints( @wp.kernel def _cfrc_ext_equality( # Model: - body_rootid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_pos: wp.array2d(dtype=wp.vec3), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_data: wp.array2d(dtype=vec11), + body_rootid: wp.array[int], + site_bodyid: wp.array[int], + site_pos: wp.array2d[wp.vec3], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], + eq_data: wp.array2d[vec11], # Data in: - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - efc_id_in: wp.array2d(dtype=int), - efc_force_in: wp.array2d(dtype=float), + xpos_in: wp.array2d[wp.vec3], + xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + efc_id_in: wp.array2d[int], + efc_force_in: wp.array2d[float], # In: - ne_connect_in: wp.array(dtype=int), - ne_weld_in: wp.array(dtype=int), + ne_connect_in: wp.array[int], + ne_weld_in: wp.array[int], # Data out: - cfrc_ext_out: wp.array2d(dtype=wp.spatial_vector), + cfrc_ext_out: wp.array2d[wp.spatial_vector], ): worldid, eqid = wp.tid() @@ -1439,22 +1439,22 @@ def transform_force(force: wp.vec3, torque: wp.vec3, offset: wp.vec3) -> wp.spat def _cfrc_ext_contact( # Model: opt_cone: int, - body_rootid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), + body_rootid: wp.array[int], + geom_bodyid: wp.array[int], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), - contact_pos_in: wp.array(dtype=wp.vec3), - contact_frame_in: wp.array(dtype=wp.mat33), - contact_friction_in: wp.array(dtype=vec5), - contact_dim_in: wp.array(dtype=int), - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - efc_force_in: wp.array2d(dtype=float), + subtree_com_in: wp.array2d[wp.vec3], + contact_pos_in: wp.array[wp.vec3], + contact_frame_in: wp.array[wp.mat33], + contact_friction_in: wp.array[vec5], + contact_dim_in: wp.array[int], + contact_geom_in: wp.array[wp.vec2i], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + efc_force_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # Data out: - cfrc_ext_out: wp.array2d(dtype=wp.spatial_vector), + cfrc_ext_out: wp.array2d[wp.spatial_vector], ): contactid = wp.tid() @@ -1584,17 +1584,17 @@ def rne_postconstraint(m: Model, d: Data): @wp.func def _accumulate_jac_dot_chain( # Model: - body_parentid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_jntid: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), + body_parentid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + jnt_type: wp.array[int], + jnt_dofadr: wp.array[int], + dof_jntid: wp.array[int], + ten_J_colind: wp.array[int], # Data in: - cdof_in: wp.array2d(dtype=wp.spatial_vector), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - cdof_dot_in: wp.array2d(dtype=wp.spatial_vector), + cdof_in: wp.array2d[wp.spatial_vector], + cvel_in: wp.array2d[wp.spatial_vector], + cdof_dot_in: wp.array2d[wp.spatial_vector], # In: offset: wp.vec3, pvel_lin: wp.vec3, @@ -1606,7 +1606,7 @@ def _accumulate_jac_dot_chain( scale: float, worldid: int, # Out: - ten_Jdot_out: wp.array2d(dtype=float), + ten_Jdot_out: wp.array2d[float], ): """Walk body chain from bodyid to root, accumulate Jdot contributions.""" ptr = rownnz - 1 @@ -1655,31 +1655,31 @@ def _accumulate_jac_dot_chain( @wp.kernel def _tendon_dot( # Model: - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_jntid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - tendon_adr: wp.array(dtype=int), - tendon_num: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_armature: wp.array2d(dtype=float), - wrap_type: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + jnt_type: wp.array[int], + jnt_dofadr: wp.array[int], + dof_jntid: wp.array[int], + site_bodyid: wp.array[int], + tendon_adr: wp.array[int], + tendon_num: wp.array[int], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_armature: wp.array2d[float], + wrap_type: wp.array[int], + wrap_objid: wp.array[int], + wrap_prm: wp.array[float], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - cdof_dot_in: wp.array2d(dtype=wp.spatial_vector), + site_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + cvel_in: wp.array2d[wp.spatial_vector], + cdof_dot_in: wp.array2d[wp.spatial_vector], # Out: - ten_Jdot_out: wp.array2d(dtype=float), + ten_Jdot_out: wp.array2d[float], ): worldid, tenid = wp.tid() @@ -1809,16 +1809,16 @@ def _tendon_dot( @wp.kernel def _tendon_bias_coef( # Model: - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_armature: wp.array2d(dtype=float), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_armature: wp.array2d[float], # Data in: - qvel_in: wp.array2d(dtype=float), + qvel_in: wp.array2d[float], # In: - ten_Jdot_in: wp.array2d(dtype=float), + ten_Jdot_in: wp.array2d[float], # Out: - ten_bias_coef_out: wp.array2d(dtype=float), + ten_bias_coef_out: wp.array2d[float], ): worldid, tenid, dofid_sparse = wp.tid() @@ -1842,16 +1842,16 @@ def _tendon_bias_coef( @wp.kernel def _tendon_bias_qfrc( # Model: - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - tendon_armature: wp.array2d(dtype=float), + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + tendon_armature: wp.array2d[float], # Data in: - ten_J_in: wp.array2d(dtype=float), + ten_J_in: wp.array2d[float], # In: - ten_bias_coef_in: wp.array2d(dtype=float), + ten_bias_coef_in: wp.array2d[float], # Out: - qfrc_out: wp.array2d(dtype=float), + qfrc_out: wp.array2d[float], ): worldid, tenid, dofid = wp.tid() @@ -1875,7 +1875,7 @@ def _tendon_bias_qfrc( @event_scope -def tendon_bias(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): +def tendon_bias(m: Model, d: Data, qfrc: wp.array2d[float]): """Add bias force due to tendon armature. Args: @@ -1933,7 +1933,7 @@ def tendon_bias(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): @wp.kernel -def _comvel_root(cvel_out: wp.array2d(dtype=wp.spatial_vector)): +def _comvel_root(cvel_out: wp.array2d[wp.spatial_vector]): worldid, elementid = wp.tid() cvel_out[worldid, 0][elementid] = 0.0 @@ -1941,19 +1941,19 @@ def _comvel_root(cvel_out: wp.array2d(dtype=wp.spatial_vector)): @wp.kernel def _comvel_branch( # Model: - body_parentid: wp.array(dtype=int), - body_jntnum: wp.array(dtype=int), - body_jntadr: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - body_branches: wp.array(dtype=int), - body_branch_start: wp.array(dtype=int), + body_parentid: wp.array[int], + body_jntnum: wp.array[int], + body_jntadr: wp.array[int], + body_dofadr: wp.array[int], + jnt_type: wp.array[int], + body_branches: wp.array[int], + body_branch_start: wp.array[int], # Data in: - qvel_in: wp.array2d(dtype=float), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + qvel_in: wp.array2d[float], + cdof_in: wp.array2d[wp.spatial_vector], # Data out: - cvel_out: wp.array2d(dtype=wp.spatial_vector), - cdof_dot_out: wp.array2d(dtype=wp.spatial_vector), + cvel_out: wp.array2d[wp.spatial_vector], + cdof_dot_out: wp.array2d[wp.spatial_vector], ): worldid, branchid = wp.tid() @@ -2042,42 +2042,42 @@ def com_vel(m: Model, d: Data): def _transmission( # Model: nv: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_bodyid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_quat: wp.array2d(dtype=wp.quat), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - actuator_trntype: wp.array(dtype=int), - actuator_trnid: wp.array(dtype=wp.vec2i), - actuator_gear: wp.array2d(dtype=wp.spatial_vector), - actuator_cranklength: wp.array2d(dtype=float), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_weldid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + jnt_type: wp.array[int], + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], + dof_bodyid: wp.array[int], + dof_parentid: wp.array[int], + site_bodyid: wp.array[int], + site_quat: wp.array2d[wp.quat], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + actuator_trntype: wp.array[int], + actuator_trnid: wp.array[wp.vec2i], + actuator_gear: wp.array2d[wp.spatial_vector], + actuator_cranklength: wp.array2d[float], # Data in: - qpos_in: wp.array2d(dtype=float), - xquat_in: wp.array2d(dtype=wp.quat), - site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - ten_J_in: wp.array2d(dtype=float), - ten_length_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], + xquat_in: wp.array2d[wp.quat], + site_xpos_in: wp.array2d[wp.vec3], + site_xmat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + ten_J_in: wp.array2d[float], + ten_length_in: wp.array2d[float], # In: - moment_nnz: wp.array(dtype=int), + moment_nnz: wp.array[int], # Data out: - actuator_length_out: wp.array2d(dtype=float), - moment_rownnz_out: wp.array2d(dtype=int), - moment_rowadr_out: wp.array2d(dtype=int), - moment_colind_out: wp.array2d(dtype=int), - actuator_moment_out: wp.array2d(dtype=float), + actuator_length_out: wp.array2d[float], + moment_rownnz_out: wp.array2d[int], + moment_rowadr_out: wp.array2d[int], + moment_colind_out: wp.array2d[int], + actuator_moment_out: wp.array2d[float], ): worldid, actid = wp.tid() trntype = actuator_trntype[actid] @@ -2448,35 +2448,35 @@ def _transmission( def _transmission_body_moment( # Model: opt_cone: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - dof_bodyid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - actuator_trnid: wp.array(dtype=wp.vec2i), - actuator_trntype_body_adr: wp.array(dtype=int), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + dof_bodyid: wp.array[int], + geom_bodyid: wp.array[int], + actuator_trnid: wp.array[wp.vec2i], + actuator_trntype_body_adr: wp.array[int], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - moment_rowadr_in: wp.array2d(dtype=int), - contact_dist_in: wp.array(dtype=float), - contact_pos_in: wp.array(dtype=wp.vec3), - contact_frame_in: wp.array(dtype=wp.mat33), - contact_includemargin_in: wp.array(dtype=float), - contact_dim_in: wp.array(dtype=int), - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - nacon_in: wp.array(dtype=int), + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + moment_rowadr_in: wp.array2d[int], + contact_dist_in: wp.array[float], + contact_pos_in: wp.array[wp.vec3], + contact_frame_in: wp.array[wp.mat33], + contact_includemargin_in: wp.array[float], + contact_dim_in: wp.array[int], + contact_geom_in: wp.array[wp.vec2i], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + nacon_in: wp.array[int], # In: efc_is_sparse: bool, # Data out: - actuator_moment_out: wp.array2d(dtype=float), + actuator_moment_out: wp.array2d[float], # Out: - actuator_trntype_body_ncon_out: wp.array2d(dtype=int), + actuator_trntype_body_ncon_out: wp.array2d[int], ): trnbodyid, conid, dofid = wp.tid() actid = actuator_trntype_body_adr[trnbodyid] @@ -2583,13 +2583,13 @@ def _transmission_body_moment( @wp.kernel def _transmission_body_moment_scale( # Model: - actuator_trntype_body_adr: wp.array(dtype=int), + actuator_trntype_body_adr: wp.array[int], # Data in: - moment_rowadr_in: wp.array2d(dtype=int), + moment_rowadr_in: wp.array2d[int], # In: - actuator_trntype_body_ncon_in: wp.array2d(dtype=int), + actuator_trntype_body_ncon_in: wp.array2d[int], # Data out: - actuator_moment_out: wp.array2d(dtype=float), + actuator_moment_out: wp.array2d[float], ): worldid, trnbodyid, dofid = wp.tid() @@ -2704,13 +2704,13 @@ def _solve_LD_sparse_fused(nv: int, nlevels: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # In: - L: wp.array3d(dtype=float), - D: wp.array2d(dtype=float), - all_updates: wp.array(dtype=wp.vec3i), - level_offsets: wp.array(dtype=int), - y: wp.array2d(dtype=float), + L: wp.array3d[float], + D: wp.array2d[float], + all_updates: wp.array[wp.vec3i], + level_offsets: wp.array[int], + y: wp.array2d[float], # Out: - x_out: wp.array2d(dtype=float), + x_out: wp.array2d[float], ): worldid, tid = wp.tid() NV = wp.static(nv) @@ -2757,10 +2757,10 @@ def _solve_LD_sparse_fused(nv: int, nlevels: int): def _solve_LD_sparse( m: Model, d: Data, - L: wp.array3d(dtype=float), - D: wp.array2d(dtype=float), - x: wp.array2d(dtype=float), - y: wp.array2d(dtype=float), + L: wp.array3d[float], + D: wp.array2d[float], + x: wp.array2d[float], + y: wp.array2d[float], ): """Computes sparse backsubstitution: x = inv(L'*D*L)*y.""" nlevels = len(m.qLD_updates) @@ -2786,11 +2786,11 @@ def _tile_cholesky_solve(tile: TileSet): @wp.kernel(module="unique", enable_backward=False) def cholesky_solve( # In: - L: wp.array3d(dtype=float), - y: wp.array2d(dtype=float), - adr: wp.array(dtype=int), + L: wp.array3d[float], + y: wp.array2d[float], + adr: wp.array[int], # Out: - x: wp.array2d(dtype=float), + x: wp.array2d[float], ): worldid, nodeid = wp.tid() TILE_SIZE = wp.static(tile.size) @@ -2804,7 +2804,7 @@ def _tile_cholesky_solve(tile: TileSet): return cholesky_solve -def _solve_LD_dense(m: Model, d: Data, L: wp.array3d(dtype=float), x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)): +def _solve_LD_dense(m: Model, d: Data, L: wp.array3d[float], x: wp.array2d[float], y: wp.array2d[float]): """Computes dense backsubstitution: x = inv(L'*L)*y.""" for tile in m.qM_tiles: wp.launch_tiled( @@ -2819,10 +2819,10 @@ def _solve_LD_dense(m: Model, d: Data, L: wp.array3d(dtype=float), x: wp.array2d def solve_LD( m: Model, d: Data, - L: wp.array3d(dtype=float), - D: wp.array2d(dtype=float), - x: wp.array2d(dtype=float), - y: wp.array2d(dtype=float), + L: wp.array3d[float], + D: wp.array2d[float], + x: wp.array2d[float], + y: wp.array2d[float], ): """Computes backsubstitution to solve a linear system of the form x = inv(L'*D*L) * y. @@ -2845,7 +2845,7 @@ def solve_LD( @event_scope -def solve_m(m: Model, d: Data, x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)): +def solve_m(m: Model, d: Data, x: wp.array2d[float], y: wp.array2d[float]): """Computes backsubstitution: x = qLD * y. Args: @@ -2864,12 +2864,12 @@ def _tile_cholesky_factorize_solve(tile: TileSet): @wp.kernel(module="unique", enable_backward=False) def cholesky_factorize_solve( # In: - M: wp.array3d(dtype=float), - y: wp.array2d(dtype=float), - adr: wp.array(dtype=int), + M: wp.array3d[float], + y: wp.array2d[float], + adr: wp.array[int], # Out: - x: wp.array2d(dtype=float), - L: wp.array3d(dtype=float), + x: wp.array2d[float], + L: wp.array3d[float], ): worldid, nodeid = wp.tid() TILE_SIZE = wp.static(tile.size) @@ -2889,10 +2889,10 @@ def _tile_cholesky_factorize_solve(tile: TileSet): def _factor_solve_i_dense( m: Model, d: Data, - M: wp.array3d(dtype=float), - x: wp.array2d(dtype=float), - y: wp.array2d(dtype=float), - L: wp.array3d(dtype=float), + M: wp.array3d[float], + x: wp.array2d[float], + y: wp.array2d[float], + L: wp.array3d[float], ): for tile in m.qM_tiles: wp.launch_tiled( @@ -2931,19 +2931,19 @@ def factor_solve_i(m, d, M, L, D, x, y): @wp.kernel def _subtree_vel_forward( # Model: - body_rootid: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_inertia: wp.array2d(dtype=wp.vec3), + body_rootid: wp.array[int], + body_mass: wp.array2d[float], + body_inertia: wp.array2d[wp.vec3], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), + xipos_in: wp.array2d[wp.vec3], + ximat_in: wp.array2d[wp.mat33], + subtree_com_in: wp.array2d[wp.vec3], + cvel_in: wp.array2d[wp.spatial_vector], # Data out: - subtree_linvel_out: wp.array2d(dtype=wp.vec3), - subtree_angmom_out: wp.array2d(dtype=wp.vec3), + subtree_linvel_out: wp.array2d[wp.vec3], + subtree_angmom_out: wp.array2d[wp.vec3], # Out: - subtree_bodyvel_out: wp.array2d(dtype=wp.spatial_vector), + subtree_bodyvel_out: wp.array2d[wp.spatial_vector], ): worldid, bodyid = wp.tid() body_mass_id = worldid % body_mass.shape[0] @@ -2971,14 +2971,14 @@ def _subtree_vel_forward( @wp.kernel def _linear_momentum( # Model: - body_parentid: wp.array(dtype=int), - body_subtreemass: wp.array2d(dtype=float), + body_parentid: wp.array[int], + body_subtreemass: wp.array2d[float], # Data in: - subtree_linvel_in: wp.array2d(dtype=wp.vec3), + subtree_linvel_in: wp.array2d[wp.vec3], # In: - body_tree_: wp.array(dtype=int), + body_tree_: wp.array[int], # Data out: - subtree_linvel_out: wp.array2d(dtype=wp.vec3), + subtree_linvel_out: wp.array2d[wp.vec3], ): worldid, nodeid = wp.tid() bodyid = body_tree_[nodeid] @@ -2991,18 +2991,18 @@ def _linear_momentum( @wp.kernel def _angular_momentum( # Model: - body_parentid: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_subtreemass: wp.array2d(dtype=float), + body_parentid: wp.array[int], + body_mass: wp.array2d[float], + body_subtreemass: wp.array2d[float], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - subtree_linvel_in: wp.array2d(dtype=wp.vec3), + xipos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + subtree_linvel_in: wp.array2d[wp.vec3], # In: - subtree_bodyvel_in: wp.array2d(dtype=wp.spatial_vector), - body_tree_: wp.array(dtype=int), + subtree_bodyvel_in: wp.array2d[wp.spatial_vector], + body_tree_: wp.array[int], # Data out: - subtree_angmom_out: wp.array2d(dtype=wp.vec3), + subtree_angmom_out: wp.array2d[wp.vec3], ): worldid, nodeid = wp.tid() bodyid = body_tree_[nodeid] @@ -3087,20 +3087,20 @@ def subtree_vel(m: Model, d: Data): @wp.kernel def _joint_tendon( # Model: - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), - tendon_jnt_adr: wp.array(dtype=int), - wrap_jnt_adr: wp.array(dtype=int), + jnt_qposadr: wp.array[int], + jnt_dofadr: wp.array[int], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + wrap_objid: wp.array[int], + wrap_prm: wp.array[float], + tendon_jnt_adr: wp.array[int], + wrap_jnt_adr: wp.array[int], # Data in: - qpos_in: wp.array2d(dtype=float), + qpos_in: wp.array2d[float], # Data out: - ten_J_out: wp.array2d(dtype=float), - ten_length_out: wp.array2d(dtype=float), + ten_J_out: wp.array2d[float], + ten_length_out: wp.array2d[float], ): worldid, wrapid = wp.tid() @@ -3126,12 +3126,12 @@ def _joint_tendon( @wp.func def _accumulate_jac_chain( # Model: - body_parentid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), + body_parentid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + ten_J_colind: wp.array[int], # Data in: - cdof_in: wp.array2d(dtype=wp.spatial_vector), + cdof_in: wp.array2d[wp.spatial_vector], # In: offset: wp.vec3, vec: wp.vec3, @@ -3141,7 +3141,7 @@ def _accumulate_jac_chain( scale: float, worldid: int, # Data out: - ten_J_out: wp.array2d(dtype=float), + ten_J_out: wp.array2d[float], ): """Walk body chain from bodyid to root, accumulate Jacobian contributions.""" ptr = rownnz - 1 @@ -3172,25 +3172,25 @@ def _accumulate_jac_chain( @wp.kernel def _spatial_site_tendon( # Model: - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - tendon_site_pair_adr: wp.array(dtype=int), - wrap_site_pair_adr: wp.array(dtype=int), - wrap_pulley_scale: wp.array(dtype=float), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + site_bodyid: wp.array[int], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + wrap_objid: wp.array[int], + tendon_site_pair_adr: wp.array[int], + wrap_site_pair_adr: wp.array[int], + wrap_pulley_scale: wp.array[float], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + site_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], # Data out: - ten_J_out: wp.array2d(dtype=float), - ten_length_out: wp.array2d(dtype=float), + ten_J_out: wp.array2d[float], + ten_length_out: wp.array2d[float], ): worldid, elementid = wp.tid() @@ -3255,33 +3255,33 @@ def _spatial_site_tendon( @wp.kernel def _spatial_geom_tendon( # Model: - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - site_bodyid: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - wrap_type: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), - tendon_geom_adr: wp.array(dtype=int), - wrap_geom_adr: wp.array(dtype=int), - wrap_pulley_scale: wp.array(dtype=float), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_dofnum: wp.array[int], + body_dofadr: wp.array[int], + geom_bodyid: wp.array[int], + geom_size: wp.array2d[wp.vec3], + site_bodyid: wp.array[int], + ten_J_rownnz: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_colind: wp.array[int], + wrap_type: wp.array[int], + wrap_objid: wp.array[int], + wrap_prm: wp.array[float], + tendon_geom_adr: wp.array[int], + wrap_geom_adr: wp.array[int], + wrap_pulley_scale: wp.array[float], # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + site_xpos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], # Data out: - ten_J_out: wp.array2d(dtype=float), - ten_length_out: wp.array2d(dtype=float), + ten_J_out: wp.array2d[float], + ten_length_out: wp.array2d[float], # Out: - wrap_geom_xpos_out: wp.array2d(dtype=wp.spatial_vector), + wrap_geom_xpos_out: wp.array2d[wp.spatial_vector], ): worldid, elementid = wp.tid() wrap_adr = wrap_geom_adr[elementid] @@ -3468,19 +3468,19 @@ def _spatial_geom_tendon( def _spatial_tendon_wrap( # Model: ntendon: int, - tendon_adr: wp.array(dtype=int), - tendon_num: wp.array(dtype=int), - wrap_type: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), + tendon_adr: wp.array[int], + tendon_num: wp.array[int], + wrap_type: wp.array[int], + wrap_objid: wp.array[int], # Data in: - site_xpos_in: wp.array2d(dtype=wp.vec3), + site_xpos_in: wp.array2d[wp.vec3], # In: - wrap_geom_xpos_in: wp.array2d(dtype=wp.spatial_vector), + wrap_geom_xpos_in: wp.array2d[wp.spatial_vector], # Data out: - ten_wrapadr_out: wp.array2d(dtype=int), - ten_wrapnum_out: wp.array2d(dtype=int), - wrap_obj_out: wp.array2d(dtype=wp.vec2i), - wrap_xpos_out: wp.array2d(dtype=wp.spatial_vector), + ten_wrapadr_out: wp.array2d[int], + ten_wrapnum_out: wp.array2d[int], + wrap_obj_out: wp.array2d[wp.vec2i], + wrap_xpos_out: wp.array2d[wp.spatial_vector], ): worldid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py index 2fabebe0..3964107d 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -38,43 +38,43 @@ _BLOCK_CHOLESKY_DIM = 32 class InverseContext: """Workspace arrays for inverse dynamics.""" - Jaref: wp.array2d(dtype=float) - search_dot: wp.array(dtype=float) - gauss: wp.array(dtype=float) - cost: wp.array(dtype=float) - prev_cost: wp.array(dtype=float) - done: wp.array(dtype=bool) - changed_efc_ids: wp.array2d(dtype=int) - changed_efc_count: wp.array(dtype=int) + Jaref: wp.array2d[float] + search_dot: wp.array[float] + gauss: wp.array[float] + cost: wp.array[float] + prev_cost: wp.array[float] + done: wp.array[bool] + changed_efc_ids: wp.array2d[int] + changed_efc_count: wp.array[int] @dataclasses.dataclass class SolverContext: """Workspace arrays for constraint solver.""" - Jaref: wp.array2d(dtype=float) - search_dot: wp.array(dtype=float) - gauss: wp.array(dtype=float) - cost: wp.array(dtype=float) - prev_cost: wp.array(dtype=float) - done: wp.array(dtype=bool) - grad: wp.array2d(dtype=float) - grad_dot: wp.array(dtype=float) - Mgrad: wp.array2d(dtype=float) - search: wp.array2d(dtype=float) - mv: wp.array2d(dtype=float) - jv: wp.array2d(dtype=float) - quad: wp.array2d(dtype=wp.vec3) - quad_gauss: wp.array(dtype=wp.vec3) - alpha: wp.array(dtype=float) - prev_grad: wp.array2d(dtype=float) - prev_Mgrad: wp.array2d(dtype=float) - beta: wp.array(dtype=float) - h: wp.array3d(dtype=float) - hfactor: wp.array3d(dtype=float) + Jaref: wp.array2d[float] + search_dot: wp.array[float] + gauss: wp.array[float] + cost: wp.array[float] + prev_cost: wp.array[float] + done: wp.array[bool] + grad: wp.array2d[float] + grad_dot: wp.array[float] + Mgrad: wp.array2d[float] + search: wp.array2d[float] + mv: wp.array2d[float] + jv: wp.array2d[float] + quad: wp.array2d[wp.vec3] + quad_gauss: wp.array[wp.vec3] + alpha: wp.array[float] + prev_grad: wp.array2d[float] + prev_Mgrad: wp.array2d[float] + beta: wp.array[float] + h: wp.array3d[float] + hfactor: wp.array3d[float] # Incremental Hessian update (Newton only) - changed_efc_ids: wp.array2d(dtype=int) - changed_efc_count: wp.array(dtype=int) + changed_efc_ids: wp.array2d[int] + changed_efc_count: wp.array[int] def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext: @@ -331,28 +331,28 @@ def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> f def linesearch_parallel_fused( # Model: opt_ls_iterations: int, - opt_impratio_invsqrt: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array[float], opt_ls_parallel_min_step: float, # Data in: - ne_in: wp.array(dtype=int), - nf_in: wp.array(dtype=int), - nefc_in: wp.array(dtype=int), - contact_friction_in: wp.array(dtype=types.vec5), - contact_efc_address_in: wp.array2d(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_D_in: wp.array2d(dtype=float), - efc_frictionloss_in: wp.array2d(dtype=float), + ne_in: wp.array[int], + nf_in: wp.array[int], + nefc_in: wp.array[int], + contact_friction_in: wp.array[types.vec5], + contact_efc_address_in: wp.array2d[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_D_in: wp.array2d[float], + efc_frictionloss_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_jv_in: wp.array2d(dtype=float), - ctx_quad_in: wp.array2d(dtype=wp.vec3), - ctx_quad_gauss_in: wp.array(dtype=wp.vec3), - ctx_done_in: wp.array(dtype=bool), + ctx_Jaref_in: wp.array2d[float], + ctx_jv_in: wp.array2d[float], + ctx_quad_in: wp.array2d[wp.vec3], + ctx_quad_gauss_in: wp.array[wp.vec3], + ctx_done_in: wp.array[bool], # Out: - cost_out: wp.array2d(dtype=float), + cost_out: wp.array2d[float], ): worldid, alphaid = wp.tid() @@ -457,10 +457,10 @@ def linesearch_parallel_best_alpha( opt_ls_iterations: int, opt_ls_parallel_min_step: float, # In: - ctx_done_in: wp.array(dtype=bool), - cost_in: wp.array2d(dtype=float), + ctx_done_in: wp.array[bool], + cost_in: wp.array2d[float], # Out: - ctx_alpha_out: wp.array(dtype=float), + ctx_alpha_out: wp.array[float], ): worldid = wp.tid() @@ -478,7 +478,7 @@ def linesearch_parallel_best_alpha( ctx_alpha_out[worldid] = _log_scale(opt_ls_parallel_min_step, 1.0, opt_ls_iterations, bestid) -def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d(dtype=float)): +def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d[float]): """Parallel linesearch with setup and teardown kernels.""" dofs_per_thread = 20 if m.nv > 50 else 50 threads_per_efc = ceil(m.nv / dofs_per_thread) @@ -575,7 +575,7 @@ def _compute_efc_eval_pt_pyramidal( nf: int, # Per-row data: efc_D: float, - efc_frictionloss: wp.array(dtype=float), + efc_frictionloss: wp.array[float], ctx_Jaref: float, ctx_jv: float, ) -> wp.vec3: @@ -607,8 +607,8 @@ def _compute_efc_eval_pt_elliptic( impratio_invsqrt: float, # Per-row data (arrays for deferred load): efc_type: int, - efc_D_in: wp.array(dtype=float), - efc_frictionloss: wp.array(dtype=float), + efc_D_in: wp.array[float], + efc_frictionloss: wp.array[float], ctx_Jaref: float, ctx_jv: float, ctx_quad: wp.vec3, @@ -652,7 +652,7 @@ def _compute_efc_eval_pt_alpha_zero_pyramidal( nf: int, # Per-row data: efc_D: float, - efc_frictionloss: wp.array(dtype=float), + efc_frictionloss: wp.array[float], ctx_Jaref: float, ctx_jv: float, ) -> wp.vec3: @@ -681,8 +681,8 @@ def _compute_efc_eval_pt_alpha_zero_elliptic( impratio_invsqrt: float, # Per-row data (arrays for deferred load): efc_type: int, - efc_D_in: wp.array(dtype=float), - efc_frictionloss: wp.array(dtype=float), + efc_D_in: wp.array[float], + efc_frictionloss: wp.array[float], ctx_Jaref: float, ctx_jv: float, ctx_quad: wp.vec3, @@ -727,7 +727,7 @@ def _compute_efc_eval_pt_3alphas_pyramidal( nf: int, # Per-row data: efc_D: float, - efc_frictionloss: wp.array(dtype=float), + efc_frictionloss: wp.array[float], ctx_Jaref: float, ctx_jv: float, ) -> tuple[wp.vec3, wp.vec3, wp.vec3]: @@ -771,8 +771,8 @@ def _compute_efc_eval_pt_3alphas_elliptic( impratio_invsqrt: float, # Per-row data (arrays for deferred load): efc_type: int, - efc_D_in: wp.array(dtype=float), - efc_frictionloss: wp.array(dtype=float), + efc_D_in: wp.array[float], + efc_frictionloss: wp.array[float], ctx_Jaref: float, ctx_jv: float, ctx_quad: wp.vec3, @@ -917,44 +917,44 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: def kernel( # Model: nv: int, - opt_tolerance: wp.array(dtype=float), - opt_ls_tolerance: wp.array(dtype=float), - opt_impratio_invsqrt: wp.array(dtype=float), - stat_meaninertia: wp.array(dtype=float), + opt_tolerance: wp.array[float], + opt_ls_tolerance: wp.array[float], + opt_impratio_invsqrt: wp.array[float], + stat_meaninertia: wp.array[float], # Data in: - ne_in: wp.array(dtype=int), - nf_in: wp.array(dtype=int), - nefc_in: wp.array(dtype=int), - qfrc_smooth_in: wp.array2d(dtype=float), - contact_friction_in: wp.array(dtype=types.vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_frictionloss_in: wp.array2d(dtype=float), + ne_in: wp.array[int], + nf_in: wp.array[int], + nefc_in: wp.array[int], + qfrc_smooth_in: wp.array2d[float], + contact_friction_in: wp.array[types.vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_frictionloss_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_search_in: wp.array2d(dtype=float), - ctx_search_dot_in: wp.array(dtype=float), - ctx_gauss_in: wp.array(dtype=float), - ctx_mv_in: wp.array2d(dtype=float), - ctx_jv_in: wp.array2d(dtype=float), - ctx_quad_in: wp.array2d(dtype=wp.vec3), - ctx_done_in: wp.array(dtype=bool), + ctx_Jaref_in: wp.array2d[float], + ctx_search_in: wp.array2d[float], + ctx_search_dot_in: wp.array[float], + ctx_gauss_in: wp.array[float], + ctx_mv_in: wp.array2d[float], + ctx_jv_in: wp.array2d[float], + ctx_quad_in: wp.array2d[wp.vec3], + ctx_done_in: wp.array[bool], # Data out: - qacc_out: wp.array2d(dtype=float), - efc_Ma_out: wp.array2d(dtype=float), + qacc_out: wp.array2d[float], + efc_Ma_out: wp.array2d[float], # Out: - ctx_Jaref_out: wp.array2d(dtype=float), - ctx_jv_out: wp.array2d(dtype=float), - ctx_quad_out: wp.array2d(dtype=wp.vec3), + ctx_Jaref_out: wp.array2d[float], + ctx_jv_out: wp.array2d[float], + ctx_quad_out: wp.array2d[wp.vec3], ): worldid, tid = wp.tid() @@ -1393,11 +1393,11 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus @wp.kernel def linesearch_zero_jv( # Data in: - nefc_in: wp.array(dtype=int), + nefc_in: wp.array[int], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_jv_out: wp.array2d(dtype=float), + ctx_jv_out: wp.array2d[float], ): worldid, efcid = wp.tid() @@ -1415,16 +1415,16 @@ def linesearch_jv_fused(is_sparse: bool, nv: int, dofs_per_thread: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # Data in: - nefc_in: wp.array(dtype=int), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), + nefc_in: wp.array[int], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], # In: - ctx_search_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_search_in: wp.array2d[float], + ctx_done_in: wp.array[bool], # Out: - ctx_jv_out: wp.array2d(dtype=float), + ctx_jv_out: wp.array2d[float], ): worldid, efcid, dofstart = wp.tid() @@ -1476,15 +1476,15 @@ def linesearch_prepare_gauss(nv: int, dofs_per_thread: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # Data in: - qfrc_smooth_in: wp.array2d(dtype=float), - efc_Ma_in: wp.array2d(dtype=float), + qfrc_smooth_in: wp.array2d[float], + efc_Ma_in: wp.array2d[float], # In: - ctx_search_in: wp.array2d(dtype=float), - ctx_gauss_in: wp.array(dtype=float), - ctx_mv_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_search_in: wp.array2d[float], + ctx_gauss_in: wp.array[float], + ctx_mv_in: wp.array2d[float], + ctx_done_in: wp.array[bool], # Out: - ctx_quad_gauss_out: wp.array(dtype=wp.vec3), + ctx_quad_gauss_out: wp.array[wp.vec3], ): worldid, dofstart = wp.tid() @@ -1523,22 +1523,22 @@ def linesearch_prepare_gauss(nv: int, dofs_per_thread: int): @wp.kernel def linesearch_prepare_quad( # Model: - opt_impratio_invsqrt: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array[float], # Data in: - nefc_in: wp.array(dtype=int), - contact_friction_in: wp.array(dtype=types.vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_D_in: wp.array2d(dtype=float), - nacon_in: wp.array(dtype=int), + nefc_in: wp.array[int], + contact_friction_in: wp.array[types.vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_D_in: wp.array2d[float], + nacon_in: wp.array[int], # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_jv_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_Jaref_in: wp.array2d[float], + ctx_jv_in: wp.array2d[float], + ctx_done_in: wp.array[bool], # Out: - ctx_quad_out: wp.array2d(dtype=wp.vec3), + ctx_quad_out: wp.array2d[wp.vec3], ): worldid, efcid = wp.tid() @@ -1619,13 +1619,13 @@ def linesearch_prepare_quad( @wp.kernel def linesearch_qacc_ma( # In: - ctx_search_in: wp.array2d(dtype=float), - ctx_mv_in: wp.array2d(dtype=float), - ctx_alpha_in: wp.array(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_search_in: wp.array2d[float], + ctx_mv_in: wp.array2d[float], + ctx_alpha_in: wp.array[float], + ctx_done_in: wp.array[bool], # Data out: - qacc_out: wp.array2d(dtype=float), - efc_Ma_out: wp.array2d(dtype=float), + qacc_out: wp.array2d[float], + efc_Ma_out: wp.array2d[float], ): worldid, dofid = wp.tid() @@ -1640,13 +1640,13 @@ def linesearch_qacc_ma( @wp.kernel def linesearch_jaref( # Data in: - nefc_in: wp.array(dtype=int), + nefc_in: wp.array[int], # In: - ctx_jv_in: wp.array2d(dtype=float), - ctx_alpha_in: wp.array(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_jv_in: wp.array2d[float], + ctx_alpha_in: wp.array[float], + ctx_done_in: wp.array[bool], # Out: - ctx_Jaref_out: wp.array2d(dtype=float), + ctx_Jaref_out: wp.array2d[float], ): worldid, efcid = wp.tid() @@ -1660,7 +1660,7 @@ def linesearch_jaref( @event_scope -def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d(dtype=float)): +def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d[float]): """Linesearch for constraint solver. Args: @@ -1706,11 +1706,11 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra @wp.kernel def solve_init_efc( # Data out: - solver_niter_out: wp.array(dtype=int), + solver_niter_out: wp.array[int], # Out: - ctx_search_dot_out: wp.array(dtype=float), - ctx_cost_out: wp.array(dtype=float), - ctx_done_out: wp.array(dtype=bool), + ctx_search_dot_out: wp.array[float], + ctx_cost_out: wp.array[float], + ctx_done_out: wp.array[bool], ): worldid = wp.tid() ctx_cost_out[worldid] = types.MJ_MAXVAL @@ -1724,15 +1724,15 @@ def solve_init_jaref(is_sparse: bool, nv: int, dofs_per_thread: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # Data in: - nefc_in: wp.array(dtype=int), - qacc_in: wp.array2d(dtype=float), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_aref_in: wp.array2d(dtype=float), + nefc_in: wp.array[int], + qacc_in: wp.array2d[float], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + efc_aref_in: wp.array2d[float], # Out: - ctx_Jaref_out: wp.array2d(dtype=float), + ctx_Jaref_out: wp.array2d[float], ): worldid, efcid, dofstart = wp.tid() @@ -1771,10 +1771,10 @@ def solve_init_jaref(is_sparse: bool, nv: int, dofs_per_thread: int): @wp.kernel def solve_init_search( # In: - ctx_Mgrad_in: wp.array2d(dtype=float), + ctx_Mgrad_in: wp.array2d[float], # Out: - ctx_search_out: wp.array2d(dtype=float), - ctx_search_dot_out: wp.array(dtype=float), + ctx_search_out: wp.array2d[float], + ctx_search_dot_out: wp.array[float], ): worldid, dofid = wp.tid() search = -1.0 * ctx_Mgrad_in[worldid, dofid] @@ -1785,12 +1785,12 @@ def solve_init_search( @wp.kernel def update_constraint_init_cost( # In: - ctx_cost_in: wp.array(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_cost_in: wp.array[float], + ctx_done_in: wp.array[bool], # Out: - ctx_gauss_out: wp.array(dtype=float), - ctx_cost_out: wp.array(dtype=float), - ctx_prev_cost_out: wp.array(dtype=float), + ctx_gauss_out: wp.array[float], + ctx_cost_out: wp.array[float], + ctx_prev_cost_out: wp.array[float], ): worldid = wp.tid() @@ -1809,29 +1809,29 @@ def update_constraint_efc(track_changes: bool): @wp.kernel(module="unique", enable_backward=False) def kernel( # Model: - opt_impratio_invsqrt: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array[float], # Data in: - ne_in: wp.array(dtype=int), - nf_in: wp.array(dtype=int), - nefc_in: wp.array(dtype=int), - contact_friction_in: wp.array(dtype=types.vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - efc_type_in: wp.array2d(dtype=int), - efc_id_in: wp.array2d(dtype=int), - efc_D_in: wp.array2d(dtype=float), - efc_frictionloss_in: wp.array2d(dtype=float), - nacon_in: wp.array(dtype=int), + ne_in: wp.array[int], + nf_in: wp.array[int], + nefc_in: wp.array[int], + contact_friction_in: wp.array[types.vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + efc_type_in: wp.array2d[int], + efc_id_in: wp.array2d[int], + efc_D_in: wp.array2d[float], + efc_frictionloss_in: wp.array2d[float], + nacon_in: wp.array[int], # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_Jaref_in: wp.array2d[float], + ctx_done_in: wp.array[bool], # Data out: - efc_force_out: wp.array2d(dtype=float), - efc_state_out: wp.array2d(dtype=int), + efc_force_out: wp.array2d[float], + efc_state_out: wp.array2d[int], # Out: - ctx_cost_out: wp.array(dtype=float), - changed_ids_out: wp.array2d(dtype=int), - changed_count_out: wp.array(dtype=int), + ctx_cost_out: wp.array[float], + changed_ids_out: wp.array2d[int], + changed_count_out: wp.array[int], ): worldid, efcid = wp.tid() @@ -1954,16 +1954,16 @@ def update_constraint_efc(track_changes: bool): @wp.kernel def update_constraint_init_qfrc_constraint_sparse( # Data in: - nefc_in: wp.array(dtype=int), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_force_in: wp.array2d(dtype=float), + nefc_in: wp.array[int], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + efc_force_in: wp.array2d[float], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Data out: - qfrc_constraint_out: wp.array2d(dtype=float), + qfrc_constraint_out: wp.array2d[float], ): worldid, efcid = wp.tid() @@ -1987,14 +1987,14 @@ def update_constraint_init_qfrc_constraint_sparse( @wp.kernel def update_constraint_init_qfrc_constraint_dense( # Data in: - nefc_in: wp.array(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_force_in: wp.array2d(dtype=float), + nefc_in: wp.array[int], + efc_J_in: wp.array3d[float], + efc_force_in: wp.array2d[float], njmax_in: int, # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Data out: - qfrc_constraint_out: wp.array2d(dtype=float), + qfrc_constraint_out: wp.array2d[float], ): worldid, dofid = wp.tid() @@ -2015,15 +2015,15 @@ def update_constraint_gauss_cost(nv: int, dofs_per_thread: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # Data in: - qacc_in: wp.array2d(dtype=float), - qfrc_smooth_in: wp.array2d(dtype=float), - qacc_smooth_in: wp.array2d(dtype=float), - efc_Ma_in: wp.array2d(dtype=float), + qacc_in: wp.array2d[float], + qfrc_smooth_in: wp.array2d[float], + qacc_smooth_in: wp.array2d[float], + efc_Ma_in: wp.array2d[float], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_gauss_out: wp.array(dtype=float), - ctx_cost_out: wp.array(dtype=float), + ctx_gauss_out: wp.array[float], + ctx_cost_out: wp.array[float], ): worldid, dofstart = wp.tid() @@ -2054,14 +2054,14 @@ def update_constraint_gauss_cost(nv: int, dofs_per_thread: int): @wp.kernel def update_gradient_h_incremental( # Data in: - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], # In: - changed_ids_in: wp.array2d(dtype=int), - changed_count_in: wp.array(dtype=int), + changed_ids_in: wp.array2d[int], + changed_count_in: wp.array[int], # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): """Incrementally update lower triangle of H for changed constraints. @@ -2101,17 +2101,17 @@ def update_gradient_h_incremental( @wp.kernel def update_gradient_h_incremental_sparse( # Data in: - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], # In: - changed_ids_in: wp.array2d(dtype=int), - changed_count_in: wp.array(dtype=int), + changed_ids_in: wp.array2d[int], + changed_count_in: wp.array[int], # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): """Incrementally update lower triangle of H for changed constraints (sparse J).""" worldid, change_idx = wp.tid() @@ -2222,9 +2222,9 @@ def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | Inver @wp.kernel def update_gradient_zero_grad_dot( # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_grad_dot_out: wp.array(dtype=float), + ctx_grad_dot_out: wp.array[float], ): worldid = wp.tid() @@ -2237,14 +2237,14 @@ def update_gradient_zero_grad_dot( @wp.kernel def update_gradient_grad( # Data in: - qfrc_smooth_in: wp.array2d(dtype=float), - qfrc_constraint_in: wp.array2d(dtype=float), - efc_Ma_in: wp.array2d(dtype=float), + qfrc_smooth_in: wp.array2d[float], + qfrc_constraint_in: wp.array2d[float], + efc_Ma_in: wp.array2d[float], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_grad_out: wp.array2d(dtype=float), - ctx_grad_dot_out: wp.array(dtype=float), + ctx_grad_out: wp.array2d[float], + ctx_grad_dot_out: wp.array[float], ): worldid, dofid = wp.tid() @@ -2259,14 +2259,14 @@ def update_gradient_grad( @wp.kernel def update_gradient_set_h_qM_lower_sparse( # Model: - qM_fullm_i: wp.array(dtype=int), - qM_fullm_j: wp.array(dtype=int), + qM_fullm_i: wp.array[int], + qM_fullm_j: wp.array[int], # Data in: - qM_in: wp.array3d(dtype=float), + qM_in: wp.array3d[float], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): worldid, elementid = wp.tid() @@ -2301,14 +2301,14 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # Data in: - nefc_in: wp.array(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + nefc_in: wp.array[int], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): worldid, elementid = wp.tid() @@ -2374,15 +2374,15 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # Data in: - nefc_in: wp.array(dtype=int), - qM_in: wp.array3d(dtype=float), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + nefc_in: wp.array[int], + qM_in: wp.array3d[float], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): worldid = wp.tid() @@ -2429,31 +2429,31 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int): @wp.kernel def update_gradient_JTCJ_sparse( # Model: - opt_impratio_invsqrt: wp.array(dtype=float), - dof_tri_row: wp.array(dtype=int), - dof_tri_col: wp.array(dtype=int), + opt_impratio_invsqrt: wp.array[float], + dof_tri_row: wp.array[int], + dof_tri_col: wp.array[int], # Data in: - contact_dist_in: wp.array(dtype=float), - contact_includemargin_in: wp.array(dtype=float), - contact_friction_in: wp.array(dtype=types.vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + contact_dist_in: wp.array[float], + contact_includemargin_in: wp.array[float], + contact_friction_in: wp.array[types.vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], naconmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_Jaref_in: wp.array2d[float], + ctx_done_in: wp.array[bool], nblocks_perblock: int, dim_block: int, # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): conid_start, elementid = wp.tid() @@ -2591,28 +2591,28 @@ def update_gradient_JTCJ_sparse( @wp.kernel def update_gradient_JTCJ_dense( # Model: - opt_impratio_invsqrt: wp.array(dtype=float), - dof_tri_row: wp.array(dtype=int), - dof_tri_col: wp.array(dtype=int), + opt_impratio_invsqrt: wp.array[float], + dof_tri_row: wp.array[int], + dof_tri_col: wp.array[int], # Data in: - contact_dist_in: wp.array(dtype=float), - contact_includemargin_in: wp.array(dtype=float), - contact_friction_in: wp.array(dtype=types.vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + contact_dist_in: wp.array[float], + contact_includemargin_in: wp.array[float], + contact_friction_in: wp.array[types.vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], naconmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_Jaref_in: wp.array2d[float], + ctx_done_in: wp.array[bool], nblocks_perblock: int, dim_block: int, # Out: - ctx_h_out: wp.array3d(dtype=float), + ctx_h_out: wp.array3d[float], ): conid_start, elementid = wp.tid() @@ -2733,11 +2733,11 @@ def update_gradient_cholesky(tile_size: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # In: - ctx_grad_in: wp.array2d(dtype=float), - h_in: wp.array3d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_grad_in: wp.array2d[float], + h_in: wp.array3d[float], + ctx_done_in: wp.array[bool], # Out: - ctx_Mgrad_out: wp.array2d(dtype=float), + ctx_Mgrad_out: wp.array2d[float], ): worldid = wp.tid() TILE_SIZE = wp.static(tile_size) @@ -2759,12 +2759,12 @@ def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int): @wp.kernel(module="unique", enable_backward=False) def kernel( # In: - ctx_done_in: wp.array(dtype=bool), - ctx_grad_in: wp.array3d(dtype=float), - ctx_h_in: wp.array3d(dtype=float), - ctx_hfactor: wp.array3d(dtype=float), + ctx_done_in: wp.array[bool], + ctx_grad_in: wp.array3d[float], + ctx_h_in: wp.array3d[float], + ctx_hfactor: wp.array3d[float], # Out: - ctx_Mgrad_out: wp.array3d(dtype=float), + ctx_Mgrad_out: wp.array3d[float], ): worldid = wp.tid() TILE_SIZE = wp.static(tile_size) @@ -2786,7 +2786,7 @@ def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int): @wp.kernel -def padding_h(nv: int, ctx_done_in: wp.array(dtype=bool), ctx_h_out: wp.array3d(dtype=float)): +def padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float]): worldid, elementid = wp.tid() if ctx_done_in[worldid]: @@ -2826,17 +2826,17 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext) @wp.kernel def _JTDAJ_sparse( # Data in: - nefc_in: wp.array(dtype=int), - efc_J_rownnz_in: wp.array2d(dtype=int), - efc_J_rowadr_in: wp.array2d(dtype=int), - efc_J_colind_in: wp.array3d(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_state_in: wp.array2d(dtype=int), + nefc_in: wp.array[int], + efc_J_rownnz_in: wp.array2d[int], + efc_J_rowadr_in: wp.array2d[int], + efc_J_colind_in: wp.array3d[int], + efc_J_in: wp.array3d[float], + efc_D_in: wp.array2d[float], + efc_state_in: wp.array2d[int], # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - h_out: wp.array3d(dtype=float), + h_out: wp.array3d[float], ): worldid, efcid = wp.tid() @@ -2870,10 +2870,10 @@ def _JTDAJ_sparse( colindj = efc_J_colind_in[worldid, 0, sparseidj] h = Ji * Jj * efc_D - wp.atomic_add(h_out[worldid, colindi], colindj, h) - - if i != j: - wp.atomic_add(h_out[worldid, colindj], colindi, h) + # Store in lower triangle only: ensure row >= col + row = wp.max(colindi, colindj) + col = wp.min(colindi, colindj) + wp.atomic_add(h_out[worldid, row], col, h) def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): @@ -3061,12 +3061,12 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte @wp.kernel def solve_prev_grad_Mgrad( # In: - ctx_grad_in: wp.array2d(dtype=float), - ctx_Mgrad_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_grad_in: wp.array2d[float], + ctx_Mgrad_in: wp.array2d[float], + ctx_done_in: wp.array[bool], # Out: - ctx_prev_grad_out: wp.array2d(dtype=float), - ctx_prev_Mgrad_out: wp.array2d(dtype=float), + ctx_prev_grad_out: wp.array2d[float], + ctx_prev_Mgrad_out: wp.array2d[float], ): worldid, dofid = wp.tid() @@ -3082,13 +3082,13 @@ def solve_beta( # Model: nv: int, # In: - ctx_grad_in: wp.array2d(dtype=float), - ctx_Mgrad_in: wp.array2d(dtype=float), - ctx_prev_grad_in: wp.array2d(dtype=float), - ctx_prev_Mgrad_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_grad_in: wp.array2d[float], + ctx_Mgrad_in: wp.array2d[float], + ctx_prev_grad_in: wp.array2d[float], + ctx_prev_Mgrad_in: wp.array2d[float], + ctx_done_in: wp.array[bool], # Out: - ctx_beta_out: wp.array(dtype=float), + ctx_beta_out: wp.array[float], ): worldid = wp.tid() @@ -3108,9 +3108,9 @@ def solve_beta( @wp.kernel def solve_zero_search_dot( # In: - ctx_done_in: wp.array(dtype=bool), + ctx_done_in: wp.array[bool], # Out: - ctx_search_dot_out: wp.array(dtype=float), + ctx_search_dot_out: wp.array[float], ): worldid = wp.tid() @@ -3125,13 +3125,13 @@ def solve_search_update( # Model: opt_solver: int, # In: - ctx_Mgrad_in: wp.array2d(dtype=float), - ctx_search_in: wp.array2d(dtype=float), - ctx_beta_in: wp.array(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_Mgrad_in: wp.array2d[float], + ctx_search_in: wp.array2d[float], + ctx_beta_in: wp.array[float], + ctx_done_in: wp.array[bool], # Out: - ctx_search_out: wp.array2d(dtype=float), - ctx_search_dot_out: wp.array(dtype=float), + ctx_search_out: wp.array2d[float], + ctx_search_dot_out: wp.array[float], ): worldid, dofid = wp.tid() @@ -3151,19 +3151,19 @@ def solve_search_update( def solve_done( # Model: nv: int, - opt_tolerance: wp.array(dtype=float), + opt_tolerance: wp.array[float], opt_iterations: int, - stat_meaninertia: wp.array(dtype=float), + stat_meaninertia: wp.array[float], # In: - ctx_grad_dot_in: wp.array(dtype=float), - ctx_cost_in: wp.array(dtype=float), - ctx_prev_cost_in: wp.array(dtype=float), - ctx_done_in: wp.array(dtype=bool), + ctx_grad_dot_in: wp.array[float], + ctx_cost_in: wp.array[float], + ctx_prev_cost_in: wp.array[float], + ctx_done_in: wp.array[bool], # Data out: - solver_niter_out: wp.array(dtype=int), + solver_niter_out: wp.array[int], # Out: - nsolving_out: wp.array(dtype=int), - ctx_done_out: wp.array(dtype=bool), + nsolving_out: wp.array[int], + ctx_done_out: wp.array[bool], ): worldid = wp.tid() @@ -3189,8 +3189,8 @@ def _solver_iteration( m: types.Model, d: types.Data, ctx: SolverContext, - step_size_cost: wp.array2d(dtype=float), - nsolving: wp.array(dtype=int), + step_size_cost: wp.array2d[float], + nsolving: wp.array[int], ): _linesearch(m, d, ctx, step_size_cost) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py index 995b2a7c..f6b3fe9d 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -69,16 +69,16 @@ def mul_m_sparse(check_skip: bool): @wp.kernel(module="unique") def _mul_m_sparse( # Model: - qM_mulm_rowadr: wp.array(dtype=int), - qM_mulm_col: wp.array(dtype=int), - qM_mulm_madr: wp.array(dtype=int), + qM_mulm_rowadr: wp.array[int], + qM_mulm_col: wp.array[int], + qM_mulm_madr: wp.array[int], # Data in: - qM_in: wp.array3d(dtype=float), + qM_in: wp.array3d[float], # In: - vec: wp.array2d(dtype=float), - skip: wp.array(dtype=bool), + vec: wp.array2d[float], + skip: wp.array[bool], # Out: - res: wp.array2d(dtype=float), + res: wp.array2d[float], ): """Sparse matmul: one thread per DOF, gather-based (no atomics).""" worldid, dofid = wp.tid() @@ -108,12 +108,12 @@ def mul_m_dense(nv: int, check_skip: bool): @wp.kernel(module="unique") def _mul_m_dense( # Data in: - qM_in: wp.array3d(dtype=float), + qM_in: wp.array3d[float], # In: - vec: wp.array2d(dtype=float), - skip: wp.array(dtype=bool), + vec: wp.array2d[float], + skip: wp.array[bool], # Out: - res: wp.array2d(dtype=float), + res: wp.array2d[float], ): worldid, i = wp.tid() @@ -133,8 +133,8 @@ def mul_m_dense(nv: int, check_skip: bool): def mul_m( m: Model, d: Data, - res: wp.array2d(dtype=float), - vec: wp.array2d(dtype=float), + res: wp.array2d[float], + vec: wp.array2d[float], skip: Optional[wp.array] = None, M: Optional[wp.array] = None, ): @@ -175,18 +175,18 @@ def mul_m( def _apply_ft( # Model: nbody: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - dof_bodyid: wp.array(dtype=int), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + dof_bodyid: wp.array[int], # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + xipos_in: wp.array2d[wp.vec3], + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], # In: - ft_in: wp.array2d(dtype=wp.spatial_vector), + ft_in: wp.array2d[wp.spatial_vector], flg_add: bool, # Out: - qfrc_out: wp.array2d(dtype=float), + qfrc_out: wp.array2d[float], ): worldid, dofid = wp.tid() cdof = cdof_in[worldid, dofid] @@ -216,7 +216,7 @@ def _apply_ft( 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( kernel=_apply_ft, 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 -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. Args: @@ -238,9 +238,7 @@ def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): @wp.func -def _decode_pyramid( - njmax_in: int, pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int -) -> wp.spatial_vector: +def _decode_pyramid(njmax_in: int, pyramid: wp.array[float], efc_address: int, mu: vec5, condim: int) -> wp.spatial_vector: """Converts pyramid representation to contact force.""" force = wp.spatial_vector() @@ -270,13 +268,13 @@ def contact_force_fn( # Model: opt_cone: int, # Data in: - contact_frame_in: wp.array(dtype=wp.mat33), - contact_friction_in: wp.array(dtype=vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - efc_force_in: wp.array2d(dtype=float), + contact_frame_in: wp.array[wp.mat33], + contact_friction_in: wp.array[vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + efc_force_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: worldid: int, contact_id: int, @@ -315,19 +313,19 @@ def contact_force_kernel( # Model: opt_cone: int, # Data in: - contact_frame_in: wp.array(dtype=wp.mat33), - contact_friction_in: wp.array(dtype=vec5), - contact_dim_in: wp.array(dtype=int), - contact_efc_address_in: wp.array2d(dtype=int), - contact_worldid_in: wp.array(dtype=int), - efc_force_in: wp.array2d(dtype=float), + contact_frame_in: wp.array[wp.mat33], + contact_friction_in: wp.array[vec5], + contact_dim_in: wp.array[int], + contact_efc_address_in: wp.array2d[int], + contact_worldid_in: wp.array[int], + efc_force_in: wp.array2d[float], njmax_in: int, - nacon_in: wp.array(dtype=int), + nacon_in: wp.array[int], # In: - contact_ids: wp.array(dtype=int), + contact_ids: wp.array[int], to_world_frame: bool, # Out: - out: wp.array(dtype=wp.spatial_vector), + out: wp.array[wp.spatial_vector], ): tid = wp.tid() @@ -353,9 +351,7 @@ def contact_force_kernel( ) -def contact_force( - m: Model, d: Data, contact_ids: wp.array(dtype=int), to_world_frame: bool, force: wp.array(dtype=wp.spatial_vector) -): +def contact_force(m: Model, d: Data, contact_ids: wp.array[int], to_world_frame: bool, force: wp.array[wp.spatial_vector]): """Compute forces for contacts in Data. Args: @@ -400,12 +396,12 @@ def transform_force(frc: wp.spatial_vector, offset: wp.vec3) -> wp.spatial_vecto @wp.func def jac_dof( # Model: - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - dof_bodyid: wp.array(dtype=int), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + dof_bodyid: wp.array[int], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], # In: point: wp.vec3, bodyid: int, @@ -441,18 +437,18 @@ def _make_jac_kernel(has_jacp: bool, has_jacr: bool): @wp.kernel(module="unique", enable_backward=False) def _jac( # Model: - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - dof_bodyid: wp.array(dtype=int), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + dof_bodyid: wp.array[int], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], # In: - point_in: wp.array(dtype=wp.vec3), - bodyid_in: wp.array(dtype=int), + point_in: wp.array[wp.vec3], + bodyid_in: wp.array[int], # Out: - jacp_out: wp.array3d(dtype=float), - jacr_out: wp.array3d(dtype=float), + jacp_out: wp.array3d[float], + jacr_out: wp.array3d[float], ): worldid, dofid = wp.tid() @@ -477,10 +473,10 @@ def _make_jac_kernel(has_jacp: bool, has_jacr: bool): def jac( m: Model, d: Data, - jacp: wp.array | None, # wp.array3d(dtype=float) - jacr: wp.array | None, # wp.array3d(dtype=float) - point: wp.array(dtype=wp.vec3), - body: wp.array(dtype=int), + jacp: wp.array | None, # wp.array3d[float] + jacr: wp.array | None, # wp.array3d[float] + point: wp.array[wp.vec3], + body: wp.array[int], ): """Compute translational and rotational Jacobian for point on body. @@ -508,17 +504,17 @@ def jac( @wp.func def jac_dot_dof( # Model: - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_bodyid: wp.array(dtype=int), - dof_jntid: wp.array(dtype=int), + body_parentid: wp.array[int], + body_rootid: wp.array[int], + jnt_type: wp.array[int], + jnt_dofadr: wp.array[int], + dof_bodyid: wp.array[int], + dof_jntid: wp.array[int], # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - cdof_dot_in: wp.array2d(dtype=wp.spatial_vector), + subtree_com_in: wp.array2d[wp.vec3], + cdof_in: wp.array2d[wp.spatial_vector], + cvel_in: wp.array2d[wp.spatial_vector], + cdof_dot_in: wp.array2d[wp.spatial_vector], # In: point: wp.vec3, bodyid: int, @@ -573,7 +569,7 @@ def jac_dot_dof( 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. 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, nmocap: int, # Data in: - time_in: wp.array(dtype=float), - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - act_in: wp.array2d(dtype=float), - qacc_warmstart_in: wp.array2d(dtype=float), - ctrl_in: wp.array2d(dtype=float), - qfrc_applied_in: wp.array2d(dtype=float), - xfrc_applied_in: wp.array2d(dtype=wp.spatial_vector), - eq_active_in: wp.array2d(dtype=bool), - mocap_pos_in: wp.array2d(dtype=wp.vec3), - mocap_quat_in: wp.array2d(dtype=wp.quat), + time_in: wp.array[float], + qpos_in: wp.array2d[float], + qvel_in: wp.array2d[float], + act_in: wp.array2d[float], + qacc_warmstart_in: wp.array2d[float], + ctrl_in: wp.array2d[float], + qfrc_applied_in: wp.array2d[float], + xfrc_applied_in: wp.array2d[wp.spatial_vector], + eq_active_in: wp.array2d[bool], + mocap_pos_in: wp.array2d[wp.vec3], + mocap_quat_in: wp.array2d[wp.quat], # In: sig_in: int, - active_in: wp.array(dtype=bool), + active_in: wp.array[bool], # Out: - state_out: wp.array2d(dtype=float), + state_out: wp.array2d[float], ): 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. 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, # In: sig_in: int, - active_in: wp.array(dtype=bool), - state_in: wp.array2d(dtype=float), + active_in: wp.array[bool], + state_in: wp.array2d[float], # Data out: - time_out: wp.array(dtype=float), - qpos_out: wp.array2d(dtype=float), - qvel_out: wp.array2d(dtype=float), - act_out: wp.array2d(dtype=float), - qacc_warmstart_out: wp.array2d(dtype=float), - ctrl_out: wp.array2d(dtype=float), - qfrc_applied_out: wp.array2d(dtype=float), - xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector), - eq_active_out: wp.array2d(dtype=bool), - mocap_pos_out: wp.array2d(dtype=wp.vec3), - mocap_quat_out: wp.array2d(dtype=wp.quat), + time_out: wp.array[float], + qpos_out: wp.array2d[float], + qvel_out: wp.array2d[float], + act_out: wp.array2d[float], + qacc_warmstart_out: wp.array2d[float], + ctrl_out: wp.array2d[float], + qfrc_applied_out: wp.array2d[float], + xfrc_applied_out: wp.array2d[wp.spatial_vector], + eq_active_out: wp.array2d[bool], + mocap_pos_out: wp.array2d[wp.vec3], + mocap_quat_out: wp.array2d[wp.quat], ): worldid = wp.tid() diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py index 11b7a7c0..7a9dc2a5 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -17,6 +17,7 @@ import enum from typing import Callable import mujoco +import numpy as np import warp as wp MJ_MINVAL = mujoco.mjMINVAL @@ -793,9 +794,18 @@ class TileSet: size: size of all the tiles in this set """ - adr: wp.array(dtype=int) + adr: wp.array[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 class Callback: @@ -933,7 +943,7 @@ class Model: geom_conaffinity: geom contact affinity (ngeom,) geom_condim: contact dimensionality (1, 3, 4, 6) (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_group: geom group inclusion/exclusion mask (ngeom,) geom_priority: geom contact priority (ngeom,) @@ -1322,7 +1332,7 @@ class Model: geom_conaffinity: array("ngeom", int) geom_condim: array("ngeom", int) geom_bodyid: array("ngeom", int) - geom_dataid: array("ngeom", int) + geom_dataid: array("*", "ngeom", int) geom_matid: array("*", "ngeom", int) geom_group: array("ngeom", int) geom_priority: array("ngeom", int) @@ -1527,70 +1537,70 @@ class Model: has_fluid: bool has_sdf_geom: bool block_dim: BlockDim - body_tree: tuple[wp.array(dtype=int), ...] - body_branches: wp.array(dtype=int) - body_branch_start: wp.array(dtype=int) + body_tree: tuple[wp.array[int], ...] + body_branches: wp.array[int] + body_branch_start: wp.array[int] mocap_bodyid: array("nmocap", int) body_fluid_ellipsoid: array("nbody", bool) - jnt_limited_slide_hinge_adr: wp.array(dtype=int) - jnt_limited_ball_adr: wp.array(dtype=int) - dof_tri_row: wp.array(dtype=int) - dof_tri_col: wp.array(dtype=int) - nxn_geom_pair: wp.array(dtype=wp.vec2i) - nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i) - nxn_pairid: wp.array(dtype=wp.vec2i) - nxn_pairid_filtered: wp.array(dtype=wp.vec2i) + jnt_limited_slide_hinge_adr: wp.array[int] + jnt_limited_ball_adr: wp.array[int] + dof_tri_row: wp.array[int] + dof_tri_col: wp.array[int] + nxn_geom_pair: wp.array[wp.vec2i] + nxn_geom_pair_filtered: wp.array[wp.vec2i] + nxn_pairid: wp.array[wp.vec2i] + nxn_pairid_filtered: wp.array[wp.vec2i] geom_pair_type_count: tuple[int, ...] geom_plugin_index: array("ngeom", int) - eq_connect_adr: wp.array(dtype=int) - eq_wld_adr: wp.array(dtype=int) - eq_jnt_adr: wp.array(dtype=int) - eq_ten_adr: wp.array(dtype=int) - eq_flex_adr: wp.array(dtype=int) - tendon_jnt_adr: wp.array(dtype=int) - tendon_site_pair_adr: wp.array(dtype=int) - tendon_geom_adr: wp.array(dtype=int) - tendon_limited_adr: wp.array(dtype=int) + eq_connect_adr: wp.array[int] + eq_wld_adr: wp.array[int] + eq_jnt_adr: wp.array[int] + eq_ten_adr: wp.array[int] + eq_flex_adr: wp.array[int] + tendon_jnt_adr: wp.array[int] + tendon_site_pair_adr: wp.array[int] + tendon_geom_adr: wp.array[int] + tendon_limited_adr: wp.array[int] max_ten_J_rownnz: int - ten_wrapadr_site: wp.array(dtype=int) - ten_wrapnum_site: wp.array(dtype=int) - wrap_jnt_adr: wp.array(dtype=int) - wrap_site_adr: wp.array(dtype=int) - wrap_site_pair_adr: wp.array(dtype=int) - wrap_geom_adr: wp.array(dtype=int) + ten_wrapadr_site: wp.array[int] + ten_wrapnum_site: wp.array[int] + wrap_jnt_adr: wp.array[int] + wrap_site_adr: wp.array[int] + wrap_site_pair_adr: wp.array[int] + wrap_geom_adr: wp.array[int] wrap_pulley_scale: array("nwrap", float) - actuator_trntype_body_adr: wp.array(dtype=int) - sensor_pos_adr: wp.array(dtype=int) - sensor_limitpos_adr: wp.array(dtype=int) - sensor_vel_adr: wp.array(dtype=int) - sensor_limitvel_adr: wp.array(dtype=int) - sensor_acc_adr: wp.array(dtype=int) - sensor_rangefinder_adr: wp.array(dtype=int) - rangefinder_sensor_adr: wp.array(dtype=int) - sensor_collision_start_adr: wp.array(dtype=int) + actuator_trntype_body_adr: wp.array[int] + sensor_pos_adr: wp.array[int] + sensor_limitpos_adr: wp.array[int] + sensor_vel_adr: wp.array[int] + sensor_limitvel_adr: wp.array[int] + sensor_acc_adr: wp.array[int] + sensor_rangefinder_adr: wp.array[int] + rangefinder_sensor_adr: wp.array[int] + sensor_collision_start_adr: wp.array[int] collision_sensor_adr: array("nsensor", int) - sensor_touch_adr: wp.array(dtype=int) - sensor_limitfrc_adr: wp.array(dtype=int) + sensor_touch_adr: wp.array[int] + sensor_limitfrc_adr: wp.array[int] sensor_e_potential: bool sensor_e_kinetic: bool - sensor_tendonactfrc_adr: wp.array(dtype=int) + sensor_tendonactfrc_adr: wp.array[int] sensor_subtree_vel: bool sensor_contact_adr: array("nsensorcontact", int) sensor_adr_to_contact_adr: array("nsensor", int) sensor_rne_postconstraint: bool sensor_rangefinder_bodyid: array("nrangefinder", int) taxel_vertadr: array("nsensortaxel", int) - taxel_sensorid: wp.array(dtype=int) + taxel_sensorid: wp.array[int] qM_tiles: tuple[TileSet, ...] - qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...] - qLD_all_updates: wp.array(dtype=wp.vec3i) - qLD_level_offsets: wp.array(dtype=int) - qM_fullm_i: wp.array(dtype=int) - qM_fullm_j: wp.array(dtype=int) + qLD_updates: tuple[wp.array[wp.vec3i], ...] + qLD_all_updates: wp.array[wp.vec3i] + qLD_level_offsets: wp.array[int] + qM_fullm_i: wp.array[int] + qM_fullm_j: wp.array[int] # 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_col: wp.array(dtype=int) # column index to gather from - qM_mulm_madr: wp.array(dtype=int) # matrix address to read + qM_mulm_rowadr: wp.array[int] # start address for each row [nv+1] + qM_mulm_col: wp.array[int] # column index to gather from + qM_mulm_madr: wp.array[int] # matrix address to read class ContactType(enum.IntFlag): @@ -1674,10 +1684,10 @@ class Constraint: type: array("nworld", "njmax", int) id: array("nworld", "njmax", int) - J_rownnz: wp.array2d(dtype=int) - J_rowadr: wp.array2d(dtype=int) - J_colind: wp.array3d(dtype=int) - J: wp.array3d(dtype=float) + J_rownnz: wp.array2d[int] + J_rowadr: wp.array2d[int] + J_colind: wp.array3d[int] + J: wp.array3d[float] pos: array("nworld", "njmax", float) margin: array("nworld", "njmax", float) D: array("nworld", "njmax_pad", float) @@ -1846,8 +1856,8 @@ class Data: moment_colind: array("nworld", "nJmom", int) actuator_moment: array("nworld", "nJmom", float) crb: array("nworld", "nbody", vec10) - qM: wp.array3d(dtype=float) - qLD: wp.array3d(dtype=float) + qM: wp.array3d[float] + qLD: wp.array3d[float] qLDiagInv: array("nworld", "nv", float) flexedge_velocity: array("nworld", "nflexedge", float) ten_velocity: array("nworld", "ntendon", float) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml index e8d98eb4..f1c4bc4a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml @@ -57,6 +57,7 @@ dev = [ "lsprotocol>=2023.0.1,<2024.0.0", "mujoco>=3.6.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 cpu = [ diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py index cf65f02a..8a825da2 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py @@ -1,4 +1,4 @@ -# Copyright 2025 The Newton Developers +# Copyright 2026 The Newton Developers # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -62,6 +62,7 @@ _OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation _KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.") _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") +_VIEWER = flags.DEFINE_enum("viewer", "mujoco", ["mujoco", "viser"], "Viewer backend (mujoco native or mjviser web)") _VIEWER_GLOBAL_STATE = {"running": True, "step_once": False} @@ -105,6 +106,70 @@ def _compile_step(m, d): 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: """Runs viewer app.""" if len(argv) < 2: @@ -115,7 +180,6 @@ def _main(argv: Sequence[str]) -> None: mjm = _load_model(epath.Path(argv[1])) mjd = mujoco.MjData(mjm) ctrls = None - ctrlid = 0 if _REPLAY.value: keys = find_keys(mjm, _REPLAY.value) 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"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...") - with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer: - opt = copy.copy(mjm.opt) + if _ENGINE.value == EngineOptions.WARP: + step_fn = _make_warp_step_fn(mjm, m, d, graph, ctrls) + else: + step_fn = _make_c_step_fn(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: - wp.capture_launch(graph) - wp.synchronize() - mjw.get_data_into(mjd, mjm, d) - - viewer.sync() - - elapsed = time.time() - start - if elapsed < mjm.opt.timestep: - time.sleep(mjm.opt.timestep - elapsed) + if _VIEWER.value == "viser": + _run_viser_viewer(mjm, mjd, step_fn) + else: + _run_passive_viewer(mjm, mjd, step_fn) def main(): diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 157315db..edde9732 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -48,34 +48,35 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) + @ffi.format_args_for_warp def _refit_bvh_shim( # Model nworld: int, - flex_dim: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_elem: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_radius: wp.array(dtype=float), - flex_shell: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_vertnum: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_type: wp.array(dtype=int), + flex_dim: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_elem: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_radius: wp.array[float], + flex_shell: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_vertadr: wp.array[int], + flex_vertnum: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + geom_type: wp.array[int], nflex: int, nflexelem: int, # Data - flexvert_xpos: wp.array2d(dtype=wp.vec3), - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), + flexvert_xpos: wp.array2d[wp.vec3], + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], # Registry rc_id: int, # Dummy output - dummy: wp.array(dtype=int), + dummy: wp.array[int], ): _m.stat = _s _m.opt = _o @@ -135,7 +136,7 @@ def _refit_bvh_jax_impl( m._impl.flex_shelldataadr, m.flex_vertadr, m.flex_vertnum, - m.geom_dataid, + jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_size, m.geom_type, m.nflex, diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index b0330c4b..530de314 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -46,67 +46,68 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) + @ffi.format_args_for_warp def _collision_shim( # Model nworld: int, block_dim: mjwp_types.BlockDim, - flex_conaffinity: wp.array(dtype=int), - flex_condim: wp.array(dtype=int), - flex_contype: wp.array(dtype=int), - flex_dim: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_friction: wp.array(dtype=wp.vec3), - flex_margin: wp.array(dtype=float), - flex_radius: wp.array(dtype=float), - flex_shell: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_shellnum: wp.array(dtype=int), - flex_vertadr: wp.array(dtype=int), - flex_vertflexid: wp.array(dtype=int), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_conaffinity: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_contype: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_gap: wp.array2d(dtype=float), - geom_margin: wp.array2d(dtype=float), + flex_conaffinity: wp.array[int], + flex_condim: wp.array[int], + flex_contype: wp.array[int], + flex_dim: wp.array[int], + flex_elem: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_friction: wp.array[wp.vec3], + flex_margin: wp.array[float], + flex_radius: wp.array[float], + flex_shell: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_shellnum: wp.array[int], + flex_vertadr: wp.array[int], + flex_vertflexid: wp.array[int], + geom_aabb: wp.array3d[wp.vec3], + geom_conaffinity: wp.array[int], + geom_condim: wp.array[int], + geom_contype: wp.array[int], + geom_dataid: wp.array2d[int], + geom_friction: wp.array2d[wp.vec3], + geom_gap: wp.array2d[float], + geom_margin: wp.array2d[float], geom_pair_type_count: tuple[int, ...], - geom_plugin_index: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_rbound: wp.array2d(dtype=float), - geom_size: wp.array2d(dtype=wp.vec3), - geom_solimp: wp.array2d(dtype=mjwp_types.vec5), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_type: wp.array(dtype=int), + geom_plugin_index: wp.array[int], + geom_priority: wp.array[int], + geom_rbound: wp.array2d[float], + geom_size: wp.array2d[wp.vec3], + geom_solimp: wp.array2d[mjwp_types.vec5], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_type: wp.array[int], has_sdf_geom: bool, - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), - hfield_ncol: wp.array(dtype=int), - hfield_nrow: wp.array(dtype=int), - hfield_size: wp.array(dtype=wp.vec4), - mesh_face: wp.array(dtype=wp.vec3i), - mesh_faceadr: wp.array(dtype=int), - mesh_graph: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_octadr: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polynum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), + hfield_adr: wp.array[int], + hfield_data: wp.array[float], + hfield_ncol: wp.array[int], + hfield_nrow: wp.array[int], + hfield_size: wp.array[wp.vec4], + mesh_face: wp.array[wp.vec3i], + mesh_faceadr: wp.array[int], + mesh_graph: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_octadr: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polymap: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polynum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_vert: wp.array[wp.vec3], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], nflex: int, nflexelem: int, nflexshelldata: int, @@ -115,25 +116,25 @@ def _collision_shim( nmaxmeshdeg: int, nmaxpolygon: int, nmeshface: int, - nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i), - nxn_pairid: wp.array(dtype=wp.vec2i), - nxn_pairid_filtered: wp.array(dtype=wp.vec2i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_child: wp.array(dtype=mjwp_types.vec8i), - oct_coeff: wp.array(dtype=mjwp_types.vec8), - pair_dim: wp.array(dtype=int), - pair_friction: wp.array2d(dtype=mjwp_types.vec5), - pair_gap: wp.array2d(dtype=float), - pair_margin: wp.array2d(dtype=float), - pair_solimp: wp.array2d(dtype=mjwp_types.vec5), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr), + nxn_geom_pair_filtered: wp.array[wp.vec2i], + nxn_pairid: wp.array[wp.vec2i], + nxn_pairid_filtered: wp.array[wp.vec2i], + oct_aabb: wp.array2d[wp.vec3], + oct_child: wp.array[mjwp_types.vec8i], + oct_coeff: wp.array[mjwp_types.vec8], + pair_dim: wp.array[int], + pair_friction: wp.array2d[mjwp_types.vec5], + pair_gap: wp.array2d[float], + pair_margin: wp.array2d[float], + pair_solimp: wp.array2d[mjwp_types.vec5], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + plugin: wp.array[int], + plugin_attr: wp.array[mjwp_types.vec_pluginattr], opt__broadphase: int, opt__broadphase_filter: int, opt__ccd_iterations: int, - opt__ccd_tolerance: wp.array(dtype=float), + opt__ccd_tolerance: wp.array[float], opt__disableflags: int, opt__enableflags: int, opt__sdf_initpoints: int, @@ -141,27 +142,27 @@ def _collision_shim( # Data naccdmax: int, naconmax: int, - flexvert_xpos: wp.array2d(dtype=wp.vec3), - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), - nacon: wp.array(dtype=int), - ncollision: wp.array(dtype=int), - contact__dim: wp.array(dtype=int), - contact__dist: wp.array(dtype=float), - contact__efc_address: wp.array2d(dtype=int), - contact__flex: wp.array(dtype=wp.vec2i), - contact__frame: wp.array(dtype=wp.mat33), - contact__friction: wp.array(dtype=mjwp_types.vec5), - contact__geom: wp.array(dtype=wp.vec2i), - contact__geomcollisionid: wp.array(dtype=int), - contact__includemargin: wp.array(dtype=float), - contact__pos: wp.array(dtype=wp.vec3), - contact__solimp: wp.array(dtype=mjwp_types.vec5), - contact__solref: wp.array(dtype=wp.vec2), - contact__solreffriction: wp.array(dtype=wp.vec2), - contact__type: wp.array(dtype=int), - contact__vert: wp.array(dtype=wp.vec2i), - contact__worldid: wp.array(dtype=int), + flexvert_xpos: wp.array2d[wp.vec3], + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], + nacon: wp.array[int], + ncollision: wp.array[int], + contact__dim: wp.array[int], + contact__dist: wp.array[float], + contact__efc_address: wp.array2d[int], + contact__flex: wp.array[wp.vec2i], + contact__frame: wp.array[wp.mat33], + contact__friction: wp.array[mjwp_types.vec5], + contact__geom: wp.array[wp.vec2i], + contact__geomcollisionid: wp.array[int], + contact__includemargin: wp.array[float], + contact__pos: wp.array[wp.vec3], + contact__solimp: wp.array[mjwp_types.vec5], + contact__solref: wp.array[wp.vec2], + contact__solreffriction: wp.array[wp.vec2], + contact__type: wp.array[int], + contact__vert: wp.array[wp.vec2i], + contact__worldid: wp.array[int], ): _m.stat = _s _m.opt = _o @@ -376,7 +377,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data): m.geom_conaffinity, m.geom_condim, m.geom_contype, - m.geom_dataid, + jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_friction, m.geom_gap, m.geom_margin, diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 3d12681a..98459fbb 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -46,210 +46,211 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) + @ffi.format_args_for_warp def _forward_shim( # Model nworld: int, - M_rowadr: wp.array(dtype=int), - M_rownnz: wp.array(dtype=int), - actuator_acc0: wp.array2d(dtype=float), - actuator_actadr: wp.array(dtype=int), - actuator_actearly: wp.array(dtype=bool), - actuator_actlimited: wp.array(dtype=bool), - actuator_actnum: wp.array(dtype=int), - actuator_actrange: wp.array2d(dtype=wp.vec2), - actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f), - actuator_biastype: wp.array(dtype=int), - actuator_cranklength: wp.array2d(dtype=float), - actuator_ctrllimited: wp.array(dtype=bool), - actuator_ctrlrange: wp.array2d(dtype=wp.vec2), - actuator_dynprm: wp.array2d(dtype=mjwp_types.vec10f), - actuator_dyntype: wp.array(dtype=int), - actuator_forcelimited: wp.array(dtype=bool), - actuator_forcerange: wp.array2d(dtype=wp.vec2), - actuator_gainprm: wp.array2d(dtype=mjwp_types.vec10f), - actuator_gaintype: wp.array(dtype=int), - actuator_gear: wp.array2d(dtype=wp.spatial_vector), - actuator_lengthrange: wp.array2d(dtype=wp.vec2), - actuator_trnid: wp.array(dtype=wp.vec2i), - actuator_trntype: wp.array(dtype=int), - actuator_trntype_body_adr: wp.array(dtype=int), + M_rowadr: wp.array[int], + M_rownnz: wp.array[int], + actuator_acc0: wp.array2d[float], + actuator_actadr: wp.array[int], + actuator_actearly: wp.array[bool], + actuator_actlimited: wp.array[bool], + actuator_actnum: wp.array[int], + actuator_actrange: wp.array2d[wp.vec2], + actuator_biasprm: wp.array2d[mjwp_types.vec10f], + actuator_biastype: wp.array[int], + actuator_cranklength: wp.array2d[float], + actuator_ctrllimited: wp.array[bool], + actuator_ctrlrange: wp.array2d[wp.vec2], + actuator_dynprm: wp.array2d[mjwp_types.vec10f], + actuator_dyntype: wp.array[int], + actuator_forcelimited: wp.array[bool], + actuator_forcerange: wp.array2d[wp.vec2], + actuator_gainprm: wp.array2d[mjwp_types.vec10f], + actuator_gaintype: wp.array[int], + actuator_gear: wp.array2d[wp.spatial_vector], + actuator_lengthrange: wp.array2d[wp.vec2], + actuator_trnid: wp.array[wp.vec2i], + actuator_trntype: wp.array[int], + actuator_trntype_body_adr: wp.array[int], block_dim: mjwp_types.BlockDim, - body_branch_start: wp.array(dtype=int), - body_branches: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_fluid_ellipsoid: wp.array(dtype=bool), - body_geomadr: wp.array(dtype=int), - body_geomnum: wp.array(dtype=int), - body_gravcomp: wp.array2d(dtype=float), - body_inertia: wp.array2d(dtype=wp.vec3), - body_invweight0: wp.array2d(dtype=wp.vec2), - body_ipos: wp.array2d(dtype=wp.vec3), - body_iquat: wp.array2d(dtype=wp.quat), - body_jntadr: wp.array(dtype=int), - body_jntnum: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_mocapid: wp.array(dtype=int), - body_parentid: wp.array(dtype=int), - body_pos: wp.array2d(dtype=wp.vec3), - body_quat: wp.array2d(dtype=wp.quat), - body_rootid: wp.array(dtype=int), - body_subtreemass: wp.array2d(dtype=float), - body_tree: tuple[wp.array(dtype=int), ...], - body_treeid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), - cam_fovy: wp.array2d(dtype=float), - cam_intrinsic: wp.array2d(dtype=wp.vec4), - cam_mat0: wp.array2d(dtype=wp.mat33), - cam_mode: wp.array(dtype=int), - cam_pos: wp.array2d(dtype=wp.vec3), - cam_pos0: wp.array2d(dtype=wp.vec3), - cam_poscom0: wp.array2d(dtype=wp.vec3), - cam_quat: wp.array2d(dtype=wp.quat), - cam_resolution: wp.array(dtype=wp.vec2i), - cam_sensorsize: wp.array(dtype=wp.vec2), - cam_targetbodyid: wp.array(dtype=int), - dof_Madr: wp.array(dtype=int), - dof_armature: wp.array2d(dtype=float), - dof_bodyid: wp.array(dtype=int), - dof_damping: wp.array2d(dtype=float), - dof_frictionloss: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), - dof_jntid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - dof_solimp: wp.array2d(dtype=mjwp_types.vec5), - dof_solref: wp.array2d(dtype=wp.vec2), - dof_treeid: wp.array(dtype=int), - dof_tri_col: wp.array(dtype=int), - dof_tri_row: wp.array(dtype=int), - eq_connect_adr: wp.array(dtype=int), - eq_data: wp.array2d(dtype=mjwp_types.vec11), - eq_flex_adr: wp.array(dtype=int), - eq_jnt_adr: wp.array(dtype=int), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_solimp: wp.array2d(dtype=mjwp_types.vec5), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_ten_adr: wp.array(dtype=int), - eq_type: wp.array(dtype=int), - eq_wld_adr: wp.array(dtype=int), - flex_bending: wp.array2d(dtype=float), - flex_centered: wp.array(dtype=bool), - flex_conaffinity: wp.array(dtype=int), - flex_condim: wp.array(dtype=int), - flex_contype: wp.array(dtype=int), - flex_damping: wp.array(dtype=float), - flex_dim: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_edgeadr: wp.array(dtype=int), - flex_edgeflap: wp.array(dtype=wp.vec2i), - flex_edgenum: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elemedge: wp.array(dtype=int), - flex_elemedgeadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_friction: wp.array(dtype=wp.vec3), - flex_margin: wp.array(dtype=float), - flex_radius: wp.array(dtype=float), - flex_shell: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_shellnum: wp.array(dtype=int), - flex_stiffness: wp.array2d(dtype=float), - flex_vert: wp.array(dtype=wp.vec3), - flex_vertadr: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), - flex_vertflexid: wp.array(dtype=int), - flex_vertnum: wp.array(dtype=int), - flexedge_J_colind: wp.array(dtype=int), - flexedge_J_rowadr: wp.array(dtype=int), - flexedge_J_rownnz: wp.array(dtype=int), - flexedge_invweight0: wp.array(dtype=float), - flexedge_length0: wp.array(dtype=float), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_bodyid: wp.array(dtype=int), - geom_conaffinity: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_contype: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_fluid: wp.array2d(dtype=float), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_gap: wp.array2d(dtype=float), - geom_group: wp.array(dtype=int), - geom_margin: wp.array2d(dtype=float), - geom_matid: wp.array2d(dtype=int), + body_branch_start: wp.array[int], + body_branches: wp.array[int], + body_dofadr: wp.array[int], + body_dofnum: wp.array[int], + body_fluid_ellipsoid: wp.array[bool], + body_geomadr: wp.array[int], + body_geomnum: wp.array[int], + body_gravcomp: wp.array2d[float], + body_inertia: wp.array2d[wp.vec3], + body_invweight0: wp.array2d[wp.vec2], + body_ipos: wp.array2d[wp.vec3], + body_iquat: wp.array2d[wp.quat], + body_jntadr: wp.array[int], + body_jntnum: wp.array[int], + body_mass: wp.array2d[float], + body_mocapid: wp.array[int], + body_parentid: wp.array[int], + body_pos: wp.array2d[wp.vec3], + body_quat: wp.array2d[wp.quat], + body_rootid: wp.array[int], + body_subtreemass: wp.array2d[float], + body_tree: tuple[wp.array[int], ...], + body_treeid: wp.array[int], + body_weldid: wp.array[int], + cam_bodyid: wp.array[int], + cam_fovy: wp.array2d[float], + cam_intrinsic: wp.array2d[wp.vec4], + cam_mat0: wp.array2d[wp.mat33], + cam_mode: wp.array[int], + cam_pos: wp.array2d[wp.vec3], + cam_pos0: wp.array2d[wp.vec3], + cam_poscom0: wp.array2d[wp.vec3], + cam_quat: wp.array2d[wp.quat], + cam_resolution: wp.array[wp.vec2i], + cam_sensorsize: wp.array[wp.vec2], + cam_targetbodyid: wp.array[int], + dof_Madr: wp.array[int], + dof_armature: wp.array2d[float], + dof_bodyid: wp.array[int], + dof_damping: wp.array2d[float], + dof_frictionloss: wp.array2d[float], + dof_invweight0: wp.array2d[float], + dof_jntid: wp.array[int], + dof_parentid: wp.array[int], + dof_solimp: wp.array2d[mjwp_types.vec5], + dof_solref: wp.array2d[wp.vec2], + dof_treeid: wp.array[int], + dof_tri_col: wp.array[int], + dof_tri_row: wp.array[int], + eq_connect_adr: wp.array[int], + eq_data: wp.array2d[mjwp_types.vec11], + eq_flex_adr: wp.array[int], + eq_jnt_adr: wp.array[int], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], + eq_solimp: wp.array2d[mjwp_types.vec5], + eq_solref: wp.array2d[wp.vec2], + eq_ten_adr: wp.array[int], + eq_type: wp.array[int], + eq_wld_adr: wp.array[int], + flex_bending: wp.array2d[float], + flex_centered: wp.array[bool], + flex_conaffinity: wp.array[int], + flex_condim: wp.array[int], + flex_contype: wp.array[int], + flex_damping: wp.array[float], + flex_dim: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_edgeadr: wp.array[int], + flex_edgeflap: wp.array[wp.vec2i], + flex_edgenum: wp.array[int], + flex_elem: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elemedge: wp.array[int], + flex_elemedgeadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_friction: wp.array[wp.vec3], + flex_margin: wp.array[float], + flex_radius: wp.array[float], + flex_shell: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_shellnum: wp.array[int], + flex_stiffness: wp.array2d[float], + flex_vert: wp.array[wp.vec3], + flex_vertadr: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_vertflexid: wp.array[int], + flex_vertnum: wp.array[int], + flexedge_J_colind: wp.array[int], + flexedge_J_rowadr: wp.array[int], + flexedge_J_rownnz: wp.array[int], + flexedge_invweight0: wp.array[float], + flexedge_length0: wp.array[float], + geom_aabb: wp.array3d[wp.vec3], + geom_bodyid: wp.array[int], + geom_conaffinity: wp.array[int], + geom_condim: wp.array[int], + geom_contype: wp.array[int], + geom_dataid: wp.array2d[int], + geom_fluid: wp.array2d[float], + geom_friction: wp.array2d[wp.vec3], + geom_gap: wp.array2d[float], + geom_group: wp.array[int], + geom_margin: wp.array2d[float], + geom_matid: wp.array2d[int], geom_pair_type_count: tuple[int, ...], - geom_plugin_index: wp.array(dtype=int), - geom_pos: wp.array2d(dtype=wp.vec3), - geom_priority: wp.array(dtype=int), - geom_quat: wp.array2d(dtype=wp.quat), - geom_rbound: wp.array2d(dtype=float), - geom_rgba: wp.array2d(dtype=wp.vec4), - geom_size: wp.array2d(dtype=wp.vec3), - geom_solimp: wp.array2d(dtype=mjwp_types.vec5), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_type: wp.array(dtype=int), + geom_plugin_index: wp.array[int], + geom_pos: wp.array2d[wp.vec3], + geom_priority: wp.array[int], + geom_quat: wp.array2d[wp.quat], + geom_rbound: wp.array2d[float], + geom_rgba: wp.array2d[wp.vec4], + geom_size: wp.array2d[wp.vec3], + geom_solimp: wp.array2d[mjwp_types.vec5], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_type: wp.array[int], has_fluid: bool, has_sdf_geom: bool, - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), - hfield_ncol: wp.array(dtype=int), - hfield_nrow: wp.array(dtype=int), - hfield_size: wp.array(dtype=wp.vec4), + hfield_adr: wp.array[int], + hfield_data: wp.array[float], + hfield_ncol: wp.array[int], + hfield_nrow: wp.array[int], + hfield_size: wp.array[wp.vec4], is_sparse: bool, - jnt_actfrclimited: wp.array(dtype=bool), - jnt_actfrcrange: wp.array2d(dtype=wp.vec2), - jnt_actgravcomp: wp.array(dtype=int), - jnt_axis: wp.array2d(dtype=wp.vec3), - jnt_bodyid: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_limited_ball_adr: wp.array(dtype=int), - jnt_limited_slide_hinge_adr: wp.array(dtype=int), - jnt_margin: wp.array2d(dtype=float), - jnt_pos: wp.array2d(dtype=wp.vec3), - jnt_qposadr: wp.array(dtype=int), - jnt_range: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array2d(dtype=mjwp_types.vec5), - jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_stiffness: wp.array2d(dtype=float), - jnt_type: wp.array(dtype=int), - light_bodyid: wp.array(dtype=int), - light_dir: wp.array2d(dtype=wp.vec3), - light_dir0: wp.array2d(dtype=wp.vec3), - light_mode: wp.array(dtype=int), - light_pos: wp.array2d(dtype=wp.vec3), - light_pos0: wp.array2d(dtype=wp.vec3), - light_poscom0: wp.array2d(dtype=wp.vec3), - light_targetbodyid: wp.array(dtype=int), - mapM2M: wp.array(dtype=int), - mat_rgba: wp.array2d(dtype=wp.vec4), + jnt_actfrclimited: wp.array[bool], + jnt_actfrcrange: wp.array2d[wp.vec2], + jnt_actgravcomp: wp.array[int], + jnt_axis: wp.array2d[wp.vec3], + jnt_bodyid: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_limited_ball_adr: wp.array[int], + jnt_limited_slide_hinge_adr: wp.array[int], + jnt_margin: wp.array2d[float], + jnt_pos: wp.array2d[wp.vec3], + jnt_qposadr: wp.array[int], + jnt_range: wp.array2d[wp.vec2], + jnt_solimp: wp.array2d[mjwp_types.vec5], + jnt_solref: wp.array2d[wp.vec2], + jnt_stiffness: wp.array2d[float], + jnt_type: wp.array[int], + light_bodyid: wp.array[int], + light_dir: wp.array2d[wp.vec3], + light_dir0: wp.array2d[wp.vec3], + light_mode: wp.array[int], + light_pos: wp.array2d[wp.vec3], + light_pos0: wp.array2d[wp.vec3], + light_poscom0: wp.array2d[wp.vec3], + light_targetbodyid: wp.array[int], + mapM2M: wp.array[int], + mat_rgba: wp.array2d[wp.vec4], max_ten_J_rownnz: int, - mesh_face: wp.array(dtype=wp.vec3i), - mesh_faceadr: wp.array(dtype=int), - mesh_graph: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_normal: wp.array(dtype=wp.vec3), - mesh_normaladr: wp.array(dtype=int), - mesh_normalnum: wp.array(dtype=int), - mesh_octadr: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polynum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_quat: wp.array(dtype=wp.quat), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), + mesh_face: wp.array[wp.vec3i], + mesh_faceadr: wp.array[int], + mesh_graph: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_normal: wp.array[wp.vec3], + mesh_normaladr: wp.array[int], + mesh_normalnum: wp.array[int], + mesh_octadr: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polymap: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polynum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_quat: wp.array[wp.quat], + mesh_vert: wp.array[wp.vec3], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], nC: int, nJten: int, na: int, @@ -283,241 +284,241 @@ def _forward_shim( nv: int, nv_pad: int, nwrap: int, - nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i), - nxn_pairid: wp.array(dtype=wp.vec2i), - nxn_pairid_filtered: wp.array(dtype=wp.vec2i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_child: wp.array(dtype=mjwp_types.vec8i), - oct_coeff: wp.array(dtype=mjwp_types.vec8), - pair_dim: wp.array(dtype=int), - pair_friction: wp.array2d(dtype=mjwp_types.vec5), - pair_gap: wp.array2d(dtype=float), - pair_margin: wp.array2d(dtype=float), - pair_solimp: wp.array2d(dtype=mjwp_types.vec5), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr), - qLD_all_updates: wp.array(dtype=wp.vec3i), - qLD_level_offsets: wp.array(dtype=int), - qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...], - qM_fullm_i: wp.array(dtype=int), - qM_fullm_j: wp.array(dtype=int), - qM_mulm_col: wp.array(dtype=int), - qM_mulm_madr: wp.array(dtype=int), - qM_mulm_rowadr: wp.array(dtype=int), + nxn_geom_pair_filtered: wp.array[wp.vec2i], + nxn_pairid: wp.array[wp.vec2i], + nxn_pairid_filtered: wp.array[wp.vec2i], + oct_aabb: wp.array2d[wp.vec3], + oct_child: wp.array[mjwp_types.vec8i], + oct_coeff: wp.array[mjwp_types.vec8], + pair_dim: wp.array[int], + pair_friction: wp.array2d[mjwp_types.vec5], + pair_gap: wp.array2d[float], + pair_margin: wp.array2d[float], + pair_solimp: wp.array2d[mjwp_types.vec5], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + plugin: wp.array[int], + plugin_attr: wp.array[mjwp_types.vec_pluginattr], + qLD_all_updates: wp.array[wp.vec3i], + qLD_level_offsets: wp.array[int], + qLD_updates: tuple[wp.array[wp.vec3i], ...], + qM_fullm_i: wp.array[int], + qM_fullm_j: wp.array[int], + qM_mulm_col: wp.array[int], + qM_mulm_madr: wp.array[int], + qM_mulm_rowadr: wp.array[int], qM_tiles: tuple[mjwp_types.TileSet, ...], - qpos0: wp.array2d(dtype=float), - qpos_spring: wp.array2d(dtype=float), - rangefinder_sensor_adr: wp.array(dtype=int), - sensor_acc_adr: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_adr_to_contact_adr: wp.array(dtype=int), - sensor_contact_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_datatype: wp.array(dtype=int), - sensor_dim: wp.array(dtype=int), + qpos0: wp.array2d[float], + qpos_spring: wp.array2d[float], + rangefinder_sensor_adr: wp.array[int], + sensor_acc_adr: wp.array[int], + sensor_adr: wp.array[int], + sensor_adr_to_contact_adr: wp.array[int], + sensor_contact_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_datatype: wp.array[int], + sensor_dim: wp.array[int], sensor_e_kinetic: bool, sensor_e_potential: bool, - sensor_intprm: wp.array2d(dtype=int), - sensor_limitfrc_adr: wp.array(dtype=int), - sensor_limitpos_adr: wp.array(dtype=int), - sensor_limitvel_adr: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_objtype: wp.array(dtype=int), - sensor_pos_adr: wp.array(dtype=int), - sensor_rangefinder_adr: wp.array(dtype=int), - sensor_rangefinder_bodyid: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_reftype: wp.array(dtype=int), + sensor_intprm: wp.array2d[int], + sensor_limitfrc_adr: wp.array[int], + sensor_limitpos_adr: wp.array[int], + sensor_limitvel_adr: wp.array[int], + sensor_objid: wp.array[int], + sensor_objtype: wp.array[int], + sensor_pos_adr: wp.array[int], + sensor_rangefinder_adr: wp.array[int], + sensor_rangefinder_bodyid: wp.array[int], + sensor_refid: wp.array[int], + sensor_reftype: wp.array[int], sensor_rne_postconstraint: bool, sensor_subtree_vel: bool, - sensor_tendonactfrc_adr: wp.array(dtype=int), - sensor_touch_adr: wp.array(dtype=int), - sensor_type: wp.array(dtype=int), - sensor_vel_adr: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_pos: wp.array2d(dtype=wp.vec3), - site_quat: wp.array2d(dtype=wp.quat), - site_size: wp.array(dtype=wp.vec3), - site_type: wp.array(dtype=int), - taxel_sensorid: wp.array(dtype=int), - taxel_vertadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - tendon_actfrclimited: wp.array(dtype=bool), - tendon_actfrcrange: wp.array2d(dtype=wp.vec2), - tendon_adr: wp.array(dtype=int), - tendon_armature: wp.array2d(dtype=float), - tendon_damping: wp.array2d(dtype=float), - tendon_frictionloss: wp.array2d(dtype=float), - tendon_geom_adr: wp.array(dtype=int), - tendon_invweight0: wp.array2d(dtype=float), - tendon_jnt_adr: wp.array(dtype=int), - tendon_length0: wp.array2d(dtype=float), - tendon_lengthspring: wp.array2d(dtype=wp.vec2), - tendon_limited_adr: wp.array(dtype=int), - tendon_margin: wp.array2d(dtype=float), - tendon_num: wp.array(dtype=int), - tendon_range: wp.array2d(dtype=wp.vec2), - tendon_site_pair_adr: wp.array(dtype=int), - tendon_solimp_fri: wp.array2d(dtype=mjwp_types.vec5), - tendon_solimp_lim: wp.array2d(dtype=mjwp_types.vec5), - tendon_solref_fri: wp.array2d(dtype=wp.vec2), - tendon_solref_lim: wp.array2d(dtype=wp.vec2), - tendon_stiffness: wp.array2d(dtype=float), - wrap_geom_adr: wp.array(dtype=int), - wrap_jnt_adr: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), - wrap_pulley_scale: wp.array(dtype=float), - wrap_site_pair_adr: wp.array(dtype=int), - wrap_type: wp.array(dtype=int), + sensor_tendonactfrc_adr: wp.array[int], + sensor_touch_adr: wp.array[int], + sensor_type: wp.array[int], + sensor_vel_adr: wp.array[int], + site_bodyid: wp.array[int], + site_pos: wp.array2d[wp.vec3], + site_quat: wp.array2d[wp.quat], + site_size: wp.array[wp.vec3], + site_type: wp.array[int], + taxel_sensorid: wp.array[int], + taxel_vertadr: wp.array[int], + ten_J_colind: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_rownnz: wp.array[int], + tendon_actfrclimited: wp.array[bool], + tendon_actfrcrange: wp.array2d[wp.vec2], + tendon_adr: wp.array[int], + tendon_armature: wp.array2d[float], + tendon_damping: wp.array2d[float], + tendon_frictionloss: wp.array2d[float], + tendon_geom_adr: wp.array[int], + tendon_invweight0: wp.array2d[float], + tendon_jnt_adr: wp.array[int], + tendon_length0: wp.array2d[float], + tendon_lengthspring: wp.array2d[wp.vec2], + tendon_limited_adr: wp.array[int], + tendon_margin: wp.array2d[float], + tendon_num: wp.array[int], + tendon_range: wp.array2d[wp.vec2], + tendon_site_pair_adr: wp.array[int], + tendon_solimp_fri: wp.array2d[mjwp_types.vec5], + tendon_solimp_lim: wp.array2d[mjwp_types.vec5], + tendon_solref_fri: wp.array2d[wp.vec2], + tendon_solref_lim: wp.array2d[wp.vec2], + tendon_stiffness: wp.array2d[float], + wrap_geom_adr: wp.array[int], + wrap_jnt_adr: wp.array[int], + wrap_objid: wp.array[int], + wrap_prm: wp.array[float], + wrap_pulley_scale: wp.array[float], + wrap_site_pair_adr: wp.array[int], + wrap_type: wp.array[int], opt__broadphase: int, opt__broadphase_filter: int, opt__ccd_iterations: int, - opt__ccd_tolerance: wp.array(dtype=float), + opt__ccd_tolerance: wp.array[float], opt__cone: int, opt__contact_sensor_maxmatch: int, - opt__density: wp.array(dtype=float), + opt__density: wp.array[float], opt__disableflags: int, opt__enableflags: int, opt__graph_conditional: bool, - opt__gravity: wp.array(dtype=wp.vec3), - opt__impratio_invsqrt: wp.array(dtype=float), + opt__gravity: wp.array[wp.vec3], + opt__impratio_invsqrt: wp.array[float], opt__iterations: int, opt__ls_iterations: int, opt__ls_parallel: bool, opt__ls_parallel_min_step: float, - opt__ls_tolerance: wp.array(dtype=float), - opt__magnetic: wp.array(dtype=wp.vec3), + opt__ls_tolerance: wp.array[float], + opt__magnetic: wp.array[wp.vec3], opt__run_collision_detection: bool, opt__sdf_initpoints: int, opt__sdf_iterations: int, opt__solver: int, - opt__timestep: wp.array(dtype=float), - opt__tolerance: wp.array(dtype=float), - opt__viscosity: wp.array(dtype=float), - opt__wind: wp.array(dtype=wp.vec3), - stat__meaninertia: wp.array(dtype=float), + opt__timestep: wp.array[float], + opt__tolerance: wp.array[float], + opt__viscosity: wp.array[float], + opt__wind: wp.array[wp.vec3], + stat__meaninertia: wp.array[float], # Data naccdmax: int, naconmax: int, njmax: int, njmax_nnz: int, - act: wp.array2d(dtype=float), - act_dot: wp.array2d(dtype=float), - actuator_force: wp.array2d(dtype=float), - actuator_length: wp.array2d(dtype=float), - actuator_moment: wp.array2d(dtype=float), - actuator_velocity: wp.array2d(dtype=float), - cacc: wp.array2d(dtype=wp.spatial_vector), - cam_xmat: wp.array2d(dtype=wp.mat33), - cam_xpos: wp.array2d(dtype=wp.vec3), - cdof: wp.array2d(dtype=wp.spatial_vector), - cdof_dot: wp.array2d(dtype=wp.spatial_vector), - cfrc_ext: wp.array2d(dtype=wp.spatial_vector), - cfrc_int: wp.array2d(dtype=wp.spatial_vector), - cinert: wp.array2d(dtype=mjwp_types.vec10), - crb: wp.array2d(dtype=mjwp_types.vec10), - ctrl: wp.array2d(dtype=float), - cvel: wp.array2d(dtype=wp.spatial_vector), - energy: wp.array(dtype=wp.vec2), - eq_active: wp.array2d(dtype=bool), - flexedge_J: wp.array2d(dtype=float), - flexedge_length: wp.array2d(dtype=float), - flexedge_velocity: wp.array2d(dtype=float), - flexvert_xpos: wp.array2d(dtype=wp.vec3), - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), - light_xdir: wp.array2d(dtype=wp.vec3), - light_xpos: wp.array2d(dtype=wp.vec3), - mocap_pos: wp.array2d(dtype=wp.vec3), - mocap_quat: wp.array2d(dtype=wp.quat), - moment_colind: wp.array2d(dtype=int), - moment_rowadr: wp.array2d(dtype=int), - moment_rownnz: wp.array2d(dtype=int), - nacon: wp.array(dtype=int), - ncollision: wp.array(dtype=int), - ne: wp.array(dtype=int), - nefc: wp.array(dtype=int), - nf: wp.array(dtype=int), - nisland: wp.array(dtype=int), - nl: wp.array(dtype=int), - qLD: wp.array3d(dtype=float), - qLDiagInv: wp.array2d(dtype=float), - qM: wp.array3d(dtype=float), - qacc: wp.array2d(dtype=float), - qacc_smooth: wp.array2d(dtype=float), - qacc_warmstart: wp.array2d(dtype=float), - qfrc_actuator: wp.array2d(dtype=float), - qfrc_applied: wp.array2d(dtype=float), - qfrc_bias: wp.array2d(dtype=float), - qfrc_constraint: wp.array2d(dtype=float), - qfrc_damper: wp.array2d(dtype=float), - qfrc_fluid: wp.array2d(dtype=float), - qfrc_gravcomp: wp.array2d(dtype=float), - qfrc_passive: wp.array2d(dtype=float), - qfrc_smooth: wp.array2d(dtype=float), - qfrc_spring: wp.array2d(dtype=float), - qpos: wp.array2d(dtype=float), - qvel: wp.array2d(dtype=float), - sensordata: wp.array2d(dtype=float), - site_xmat: wp.array2d(dtype=wp.mat33), - site_xpos: wp.array2d(dtype=wp.vec3), - solver_niter: wp.array(dtype=int), - subtree_angmom: wp.array2d(dtype=wp.vec3), - subtree_com: wp.array2d(dtype=wp.vec3), - subtree_linvel: wp.array2d(dtype=wp.vec3), - ten_J: wp.array2d(dtype=float), - ten_length: wp.array2d(dtype=float), - ten_velocity: wp.array2d(dtype=float), - ten_wrapadr: wp.array2d(dtype=int), - ten_wrapnum: wp.array2d(dtype=int), - time: wp.array(dtype=float), - tree_island: wp.array2d(dtype=int), - wrap_obj: wp.array2d(dtype=wp.vec2i), - wrap_xpos: wp.array2d(dtype=wp.spatial_vector), - xanchor: wp.array2d(dtype=wp.vec3), - xaxis: wp.array2d(dtype=wp.vec3), - xfrc_applied: wp.array2d(dtype=wp.spatial_vector), - ximat: wp.array2d(dtype=wp.mat33), - xipos: wp.array2d(dtype=wp.vec3), - xmat: wp.array2d(dtype=wp.mat33), - xpos: wp.array2d(dtype=wp.vec3), - xquat: wp.array2d(dtype=wp.quat), - contact__dim: wp.array(dtype=int), - contact__dist: wp.array(dtype=float), - contact__efc_address: wp.array2d(dtype=int), - contact__flex: wp.array(dtype=wp.vec2i), - contact__frame: wp.array(dtype=wp.mat33), - contact__friction: wp.array(dtype=mjwp_types.vec5), - contact__geom: wp.array(dtype=wp.vec2i), - contact__geomcollisionid: wp.array(dtype=int), - contact__includemargin: wp.array(dtype=float), - contact__pos: wp.array(dtype=wp.vec3), - contact__solimp: wp.array(dtype=mjwp_types.vec5), - contact__solref: wp.array(dtype=wp.vec2), - contact__solreffriction: wp.array(dtype=wp.vec2), - contact__type: wp.array(dtype=int), - contact__vert: wp.array(dtype=wp.vec2i), - contact__worldid: wp.array(dtype=int), - efc__D: wp.array2d(dtype=float), - efc__J: wp.array3d(dtype=float), - efc__J_colind: wp.array3d(dtype=int), - efc__J_rowadr: wp.array2d(dtype=int), - efc__J_rownnz: wp.array2d(dtype=int), - efc__Ma: wp.array2d(dtype=float), - efc__aref: wp.array2d(dtype=float), - efc__force: wp.array2d(dtype=float), - efc__frictionloss: wp.array2d(dtype=float), - efc__id: wp.array2d(dtype=int), - efc__margin: wp.array2d(dtype=float), - efc__pos: wp.array2d(dtype=float), - efc__state: wp.array2d(dtype=int), - efc__type: wp.array2d(dtype=int), - efc__vel: wp.array2d(dtype=float), + act: wp.array2d[float], + act_dot: wp.array2d[float], + actuator_force: wp.array2d[float], + actuator_length: wp.array2d[float], + actuator_moment: wp.array2d[float], + actuator_velocity: wp.array2d[float], + cacc: wp.array2d[wp.spatial_vector], + cam_xmat: wp.array2d[wp.mat33], + cam_xpos: wp.array2d[wp.vec3], + cdof: wp.array2d[wp.spatial_vector], + cdof_dot: wp.array2d[wp.spatial_vector], + cfrc_ext: wp.array2d[wp.spatial_vector], + cfrc_int: wp.array2d[wp.spatial_vector], + cinert: wp.array2d[mjwp_types.vec10], + crb: wp.array2d[mjwp_types.vec10], + ctrl: wp.array2d[float], + cvel: wp.array2d[wp.spatial_vector], + energy: wp.array[wp.vec2], + eq_active: wp.array2d[bool], + flexedge_J: wp.array2d[float], + flexedge_length: wp.array2d[float], + flexedge_velocity: wp.array2d[float], + flexvert_xpos: wp.array2d[wp.vec3], + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], + light_xdir: wp.array2d[wp.vec3], + light_xpos: wp.array2d[wp.vec3], + mocap_pos: wp.array2d[wp.vec3], + mocap_quat: wp.array2d[wp.quat], + moment_colind: wp.array2d[int], + moment_rowadr: wp.array2d[int], + moment_rownnz: wp.array2d[int], + nacon: wp.array[int], + ncollision: wp.array[int], + ne: wp.array[int], + nefc: wp.array[int], + nf: wp.array[int], + nisland: wp.array[int], + nl: wp.array[int], + qLD: wp.array3d[float], + qLDiagInv: wp.array2d[float], + qM: wp.array3d[float], + qacc: wp.array2d[float], + qacc_smooth: wp.array2d[float], + qacc_warmstart: wp.array2d[float], + qfrc_actuator: wp.array2d[float], + qfrc_applied: wp.array2d[float], + qfrc_bias: wp.array2d[float], + qfrc_constraint: wp.array2d[float], + qfrc_damper: wp.array2d[float], + qfrc_fluid: wp.array2d[float], + qfrc_gravcomp: wp.array2d[float], + qfrc_passive: wp.array2d[float], + qfrc_smooth: wp.array2d[float], + qfrc_spring: wp.array2d[float], + qpos: wp.array2d[float], + qvel: wp.array2d[float], + sensordata: wp.array2d[float], + site_xmat: wp.array2d[wp.mat33], + site_xpos: wp.array2d[wp.vec3], + solver_niter: wp.array[int], + subtree_angmom: wp.array2d[wp.vec3], + subtree_com: wp.array2d[wp.vec3], + subtree_linvel: wp.array2d[wp.vec3], + ten_J: wp.array2d[float], + ten_length: wp.array2d[float], + ten_velocity: wp.array2d[float], + ten_wrapadr: wp.array2d[int], + ten_wrapnum: wp.array2d[int], + time: wp.array[float], + tree_island: wp.array2d[int], + wrap_obj: wp.array2d[wp.vec2i], + wrap_xpos: wp.array2d[wp.spatial_vector], + xanchor: wp.array2d[wp.vec3], + xaxis: wp.array2d[wp.vec3], + xfrc_applied: wp.array2d[wp.spatial_vector], + ximat: wp.array2d[wp.mat33], + xipos: wp.array2d[wp.vec3], + xmat: wp.array2d[wp.mat33], + xpos: wp.array2d[wp.vec3], + xquat: wp.array2d[wp.quat], + contact__dim: wp.array[int], + contact__dist: wp.array[float], + contact__efc_address: wp.array2d[int], + contact__flex: wp.array[wp.vec2i], + contact__frame: wp.array[wp.mat33], + contact__friction: wp.array[mjwp_types.vec5], + contact__geom: wp.array[wp.vec2i], + contact__geomcollisionid: wp.array[int], + contact__includemargin: wp.array[float], + contact__pos: wp.array[wp.vec3], + contact__solimp: wp.array[mjwp_types.vec5], + contact__solref: wp.array[wp.vec2], + contact__solreffriction: wp.array[wp.vec2], + contact__type: wp.array[int], + contact__vert: wp.array[wp.vec2i], + contact__worldid: wp.array[int], + efc__D: wp.array2d[float], + efc__J: wp.array3d[float], + efc__J_colind: wp.array3d[int], + efc__J_rowadr: wp.array2d[int], + efc__J_rownnz: wp.array2d[int], + efc__Ma: wp.array2d[float], + efc__aref: wp.array2d[float], + efc__force: wp.array2d[float], + efc__frictionloss: wp.array2d[float], + efc__id: wp.array2d[int], + efc__margin: wp.array2d[float], + efc__pos: wp.array2d[float], + efc__state: wp.array2d[int], + efc__type: wp.array2d[int], + efc__vel: wp.array2d[float], ): _m.stat = _s _m.opt = _o @@ -1512,7 +1513,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.geom_conaffinity, m.geom_condim, m.geom_contype, - m.geom_dataid, + jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_fluid, m.geom_friction, m.geom_gap, @@ -1979,206 +1980,206 @@ def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data): def _step_shim( # Model nworld: int, - M_rowadr: wp.array(dtype=int), - M_rownnz: wp.array(dtype=int), - actuator_acc0: wp.array2d(dtype=float), - actuator_actadr: wp.array(dtype=int), - actuator_actearly: wp.array(dtype=bool), - actuator_actlimited: wp.array(dtype=bool), - actuator_actnum: wp.array(dtype=int), - actuator_actrange: wp.array2d(dtype=wp.vec2), - actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f), - actuator_biastype: wp.array(dtype=int), - actuator_cranklength: wp.array2d(dtype=float), - actuator_ctrllimited: wp.array(dtype=bool), - actuator_ctrlrange: wp.array2d(dtype=wp.vec2), - actuator_dynprm: wp.array2d(dtype=mjwp_types.vec10f), - actuator_dyntype: wp.array(dtype=int), - actuator_forcelimited: wp.array(dtype=bool), - actuator_forcerange: wp.array2d(dtype=wp.vec2), - actuator_gainprm: wp.array2d(dtype=mjwp_types.vec10f), - actuator_gaintype: wp.array(dtype=int), - actuator_gear: wp.array2d(dtype=wp.spatial_vector), - actuator_lengthrange: wp.array2d(dtype=wp.vec2), - actuator_trnid: wp.array(dtype=wp.vec2i), - actuator_trntype: wp.array(dtype=int), - actuator_trntype_body_adr: wp.array(dtype=int), + M_rowadr: wp.array[int], + M_rownnz: wp.array[int], + actuator_acc0: wp.array2d[float], + actuator_actadr: wp.array[int], + actuator_actearly: wp.array[bool], + actuator_actlimited: wp.array[bool], + actuator_actnum: wp.array[int], + actuator_actrange: wp.array2d[wp.vec2], + actuator_biasprm: wp.array2d[mjwp_types.vec10f], + actuator_biastype: wp.array[int], + actuator_cranklength: wp.array2d[float], + actuator_ctrllimited: wp.array[bool], + actuator_ctrlrange: wp.array2d[wp.vec2], + actuator_dynprm: wp.array2d[mjwp_types.vec10f], + actuator_dyntype: wp.array[int], + actuator_forcelimited: wp.array[bool], + actuator_forcerange: wp.array2d[wp.vec2], + actuator_gainprm: wp.array2d[mjwp_types.vec10f], + actuator_gaintype: wp.array[int], + actuator_gear: wp.array2d[wp.spatial_vector], + actuator_lengthrange: wp.array2d[wp.vec2], + actuator_trnid: wp.array[wp.vec2i], + actuator_trntype: wp.array[int], + actuator_trntype_body_adr: wp.array[int], block_dim: mjwp_types.BlockDim, - body_branch_start: wp.array(dtype=int), - body_branches: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_fluid_ellipsoid: wp.array(dtype=bool), - body_geomadr: wp.array(dtype=int), - body_geomnum: wp.array(dtype=int), - body_gravcomp: wp.array2d(dtype=float), - body_inertia: wp.array2d(dtype=wp.vec3), - body_invweight0: wp.array2d(dtype=wp.vec2), - body_ipos: wp.array2d(dtype=wp.vec3), - body_iquat: wp.array2d(dtype=wp.quat), - body_jntadr: wp.array(dtype=int), - body_jntnum: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_mocapid: wp.array(dtype=int), - body_parentid: wp.array(dtype=int), - body_pos: wp.array2d(dtype=wp.vec3), - body_quat: wp.array2d(dtype=wp.quat), - body_rootid: wp.array(dtype=int), - body_subtreemass: wp.array2d(dtype=float), - body_tree: tuple[wp.array(dtype=int), ...], - body_treeid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - cam_bodyid: wp.array(dtype=int), - cam_fovy: wp.array2d(dtype=float), - cam_intrinsic: wp.array2d(dtype=wp.vec4), - cam_mat0: wp.array2d(dtype=wp.mat33), - cam_mode: wp.array(dtype=int), - cam_pos: wp.array2d(dtype=wp.vec3), - cam_pos0: wp.array2d(dtype=wp.vec3), - cam_poscom0: wp.array2d(dtype=wp.vec3), - cam_quat: wp.array2d(dtype=wp.quat), - cam_resolution: wp.array(dtype=wp.vec2i), - cam_sensorsize: wp.array(dtype=wp.vec2), - cam_targetbodyid: wp.array(dtype=int), - dof_Madr: wp.array(dtype=int), - dof_armature: wp.array2d(dtype=float), - dof_bodyid: wp.array(dtype=int), - dof_damping: wp.array2d(dtype=float), - dof_frictionloss: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), - dof_jntid: wp.array(dtype=int), - dof_parentid: wp.array(dtype=int), - dof_solimp: wp.array2d(dtype=mjwp_types.vec5), - dof_solref: wp.array2d(dtype=wp.vec2), - dof_treeid: wp.array(dtype=int), - dof_tri_col: wp.array(dtype=int), - dof_tri_row: wp.array(dtype=int), - eq_connect_adr: wp.array(dtype=int), - eq_data: wp.array2d(dtype=mjwp_types.vec11), - eq_flex_adr: wp.array(dtype=int), - eq_jnt_adr: wp.array(dtype=int), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_solimp: wp.array2d(dtype=mjwp_types.vec5), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_ten_adr: wp.array(dtype=int), - eq_type: wp.array(dtype=int), - eq_wld_adr: wp.array(dtype=int), - flex_bending: wp.array2d(dtype=float), - flex_centered: wp.array(dtype=bool), - flex_conaffinity: wp.array(dtype=int), - flex_condim: wp.array(dtype=int), - flex_contype: wp.array(dtype=int), - flex_damping: wp.array(dtype=float), - flex_dim: wp.array(dtype=int), - flex_edge: wp.array(dtype=wp.vec2i), - flex_edgeadr: wp.array(dtype=int), - flex_edgeflap: wp.array(dtype=wp.vec2i), - flex_edgenum: wp.array(dtype=int), - flex_elem: wp.array(dtype=int), - flex_elemadr: wp.array(dtype=int), - flex_elemdataadr: wp.array(dtype=int), - flex_elemedge: wp.array(dtype=int), - flex_elemedgeadr: wp.array(dtype=int), - flex_elemnum: wp.array(dtype=int), - flex_friction: wp.array(dtype=wp.vec3), - flex_margin: wp.array(dtype=float), - flex_radius: wp.array(dtype=float), - flex_shell: wp.array(dtype=int), - flex_shelldataadr: wp.array(dtype=int), - flex_shellnum: wp.array(dtype=int), - flex_stiffness: wp.array2d(dtype=float), - flex_vert: wp.array(dtype=wp.vec3), - flex_vertadr: wp.array(dtype=int), - flex_vertbodyid: wp.array(dtype=int), - flex_vertflexid: wp.array(dtype=int), - flex_vertnum: wp.array(dtype=int), - flexedge_J_colind: wp.array(dtype=int), - flexedge_J_rowadr: wp.array(dtype=int), - flexedge_J_rownnz: wp.array(dtype=int), - flexedge_invweight0: wp.array(dtype=float), - flexedge_length0: wp.array(dtype=float), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_bodyid: wp.array(dtype=int), - geom_conaffinity: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_contype: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_fluid: wp.array2d(dtype=float), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_gap: wp.array2d(dtype=float), - geom_group: wp.array(dtype=int), - geom_margin: wp.array2d(dtype=float), - geom_matid: wp.array2d(dtype=int), + body_branch_start: wp.array[int], + body_branches: wp.array[int], + body_dofadr: wp.array[int], + body_dofnum: wp.array[int], + body_fluid_ellipsoid: wp.array[bool], + body_geomadr: wp.array[int], + body_geomnum: wp.array[int], + body_gravcomp: wp.array2d[float], + body_inertia: wp.array2d[wp.vec3], + body_invweight0: wp.array2d[wp.vec2], + body_ipos: wp.array2d[wp.vec3], + body_iquat: wp.array2d[wp.quat], + body_jntadr: wp.array[int], + body_jntnum: wp.array[int], + body_mass: wp.array2d[float], + body_mocapid: wp.array[int], + body_parentid: wp.array[int], + body_pos: wp.array2d[wp.vec3], + body_quat: wp.array2d[wp.quat], + body_rootid: wp.array[int], + body_subtreemass: wp.array2d[float], + body_tree: tuple[wp.array[int], ...], + body_treeid: wp.array[int], + body_weldid: wp.array[int], + cam_bodyid: wp.array[int], + cam_fovy: wp.array2d[float], + cam_intrinsic: wp.array2d[wp.vec4], + cam_mat0: wp.array2d[wp.mat33], + cam_mode: wp.array[int], + cam_pos: wp.array2d[wp.vec3], + cam_pos0: wp.array2d[wp.vec3], + cam_poscom0: wp.array2d[wp.vec3], + cam_quat: wp.array2d[wp.quat], + cam_resolution: wp.array[wp.vec2i], + cam_sensorsize: wp.array[wp.vec2], + cam_targetbodyid: wp.array[int], + dof_Madr: wp.array[int], + dof_armature: wp.array2d[float], + dof_bodyid: wp.array[int], + dof_damping: wp.array2d[float], + dof_frictionloss: wp.array2d[float], + dof_invweight0: wp.array2d[float], + dof_jntid: wp.array[int], + dof_parentid: wp.array[int], + dof_solimp: wp.array2d[mjwp_types.vec5], + dof_solref: wp.array2d[wp.vec2], + dof_treeid: wp.array[int], + dof_tri_col: wp.array[int], + dof_tri_row: wp.array[int], + eq_connect_adr: wp.array[int], + eq_data: wp.array2d[mjwp_types.vec11], + eq_flex_adr: wp.array[int], + eq_jnt_adr: wp.array[int], + eq_obj1id: wp.array[int], + eq_obj2id: wp.array[int], + eq_objtype: wp.array[int], + eq_solimp: wp.array2d[mjwp_types.vec5], + eq_solref: wp.array2d[wp.vec2], + eq_ten_adr: wp.array[int], + eq_type: wp.array[int], + eq_wld_adr: wp.array[int], + flex_bending: wp.array2d[float], + flex_centered: wp.array[bool], + flex_conaffinity: wp.array[int], + flex_condim: wp.array[int], + flex_contype: wp.array[int], + flex_damping: wp.array[float], + flex_dim: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_edgeadr: wp.array[int], + flex_edgeflap: wp.array[wp.vec2i], + flex_edgenum: wp.array[int], + flex_elem: wp.array[int], + flex_elemadr: wp.array[int], + flex_elemdataadr: wp.array[int], + flex_elemedge: wp.array[int], + flex_elemedgeadr: wp.array[int], + flex_elemnum: wp.array[int], + flex_friction: wp.array[wp.vec3], + flex_margin: wp.array[float], + flex_radius: wp.array[float], + flex_shell: wp.array[int], + flex_shelldataadr: wp.array[int], + flex_shellnum: wp.array[int], + flex_stiffness: wp.array2d[float], + flex_vert: wp.array[wp.vec3], + flex_vertadr: wp.array[int], + flex_vertbodyid: wp.array[int], + flex_vertflexid: wp.array[int], + flex_vertnum: wp.array[int], + flexedge_J_colind: wp.array[int], + flexedge_J_rowadr: wp.array[int], + flexedge_J_rownnz: wp.array[int], + flexedge_invweight0: wp.array[float], + flexedge_length0: wp.array[float], + geom_aabb: wp.array3d[wp.vec3], + geom_bodyid: wp.array[int], + geom_conaffinity: wp.array[int], + geom_condim: wp.array[int], + geom_contype: wp.array[int], + geom_dataid: wp.array2d[int], + geom_fluid: wp.array2d[float], + geom_friction: wp.array2d[wp.vec3], + geom_gap: wp.array2d[float], + geom_group: wp.array[int], + geom_margin: wp.array2d[float], + geom_matid: wp.array2d[int], geom_pair_type_count: tuple[int, ...], - geom_plugin_index: wp.array(dtype=int), - geom_pos: wp.array2d(dtype=wp.vec3), - geom_priority: wp.array(dtype=int), - geom_quat: wp.array2d(dtype=wp.quat), - geom_rbound: wp.array2d(dtype=float), - geom_rgba: wp.array2d(dtype=wp.vec4), - geom_size: wp.array2d(dtype=wp.vec3), - geom_solimp: wp.array2d(dtype=mjwp_types.vec5), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_type: wp.array(dtype=int), + geom_plugin_index: wp.array[int], + geom_pos: wp.array2d[wp.vec3], + geom_priority: wp.array[int], + geom_quat: wp.array2d[wp.quat], + geom_rbound: wp.array2d[float], + geom_rgba: wp.array2d[wp.vec4], + geom_size: wp.array2d[wp.vec3], + geom_solimp: wp.array2d[mjwp_types.vec5], + geom_solmix: wp.array2d[float], + geom_solref: wp.array2d[wp.vec2], + geom_type: wp.array[int], has_fluid: bool, has_sdf_geom: bool, - hfield_adr: wp.array(dtype=int), - hfield_data: wp.array(dtype=float), - hfield_ncol: wp.array(dtype=int), - hfield_nrow: wp.array(dtype=int), - hfield_size: wp.array(dtype=wp.vec4), + hfield_adr: wp.array[int], + hfield_data: wp.array[float], + hfield_ncol: wp.array[int], + hfield_nrow: wp.array[int], + hfield_size: wp.array[wp.vec4], is_sparse: bool, - jnt_actfrclimited: wp.array(dtype=bool), - jnt_actfrcrange: wp.array2d(dtype=wp.vec2), - jnt_actgravcomp: wp.array(dtype=int), - jnt_axis: wp.array2d(dtype=wp.vec3), - jnt_bodyid: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_limited_ball_adr: wp.array(dtype=int), - jnt_limited_slide_hinge_adr: wp.array(dtype=int), - jnt_margin: wp.array2d(dtype=float), - jnt_pos: wp.array2d(dtype=wp.vec3), - jnt_qposadr: wp.array(dtype=int), - jnt_range: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array2d(dtype=mjwp_types.vec5), - jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_stiffness: wp.array2d(dtype=float), - jnt_type: wp.array(dtype=int), - light_bodyid: wp.array(dtype=int), - light_dir: wp.array2d(dtype=wp.vec3), - light_dir0: wp.array2d(dtype=wp.vec3), - light_mode: wp.array(dtype=int), - light_pos: wp.array2d(dtype=wp.vec3), - light_pos0: wp.array2d(dtype=wp.vec3), - light_poscom0: wp.array2d(dtype=wp.vec3), - light_targetbodyid: wp.array(dtype=int), - mapM2M: wp.array(dtype=int), - mat_rgba: wp.array2d(dtype=wp.vec4), + jnt_actfrclimited: wp.array[bool], + jnt_actfrcrange: wp.array2d[wp.vec2], + jnt_actgravcomp: wp.array[int], + jnt_axis: wp.array2d[wp.vec3], + jnt_bodyid: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_limited_ball_adr: wp.array[int], + jnt_limited_slide_hinge_adr: wp.array[int], + jnt_margin: wp.array2d[float], + jnt_pos: wp.array2d[wp.vec3], + jnt_qposadr: wp.array[int], + jnt_range: wp.array2d[wp.vec2], + jnt_solimp: wp.array2d[mjwp_types.vec5], + jnt_solref: wp.array2d[wp.vec2], + jnt_stiffness: wp.array2d[float], + jnt_type: wp.array[int], + light_bodyid: wp.array[int], + light_dir: wp.array2d[wp.vec3], + light_dir0: wp.array2d[wp.vec3], + light_mode: wp.array[int], + light_pos: wp.array2d[wp.vec3], + light_pos0: wp.array2d[wp.vec3], + light_poscom0: wp.array2d[wp.vec3], + light_targetbodyid: wp.array[int], + mapM2M: wp.array[int], + mat_rgba: wp.array2d[wp.vec4], max_ten_J_rownnz: int, - mesh_face: wp.array(dtype=wp.vec3i), - mesh_faceadr: wp.array(dtype=int), - mesh_graph: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_normal: wp.array(dtype=wp.vec3), - mesh_normaladr: wp.array(dtype=int), - mesh_normalnum: wp.array(dtype=int), - mesh_octadr: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polynum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_quat: wp.array(dtype=wp.quat), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), + mesh_face: wp.array[wp.vec3i], + mesh_faceadr: wp.array[int], + mesh_graph: wp.array[int], + mesh_graphadr: wp.array[int], + mesh_normal: wp.array[wp.vec3], + mesh_normaladr: wp.array[int], + mesh_normalnum: wp.array[int], + mesh_octadr: wp.array[int], + mesh_polyadr: wp.array[int], + mesh_polymap: wp.array[int], + mesh_polymapadr: wp.array[int], + mesh_polymapnum: wp.array[int], + mesh_polynormal: wp.array[wp.vec3], + mesh_polynum: wp.array[int], + mesh_polyvert: wp.array[int], + mesh_polyvertadr: wp.array[int], + mesh_polyvertnum: wp.array[int], + mesh_quat: wp.array[wp.quat], + mesh_vert: wp.array[wp.vec3], + mesh_vertadr: wp.array[int], + mesh_vertnum: wp.array[int], nC: int, nJten: int, nM: int, @@ -2213,242 +2214,242 @@ def _step_shim( nv: int, nv_pad: int, nwrap: int, - nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i), - nxn_pairid: wp.array(dtype=wp.vec2i), - nxn_pairid_filtered: wp.array(dtype=wp.vec2i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_child: wp.array(dtype=mjwp_types.vec8i), - oct_coeff: wp.array(dtype=mjwp_types.vec8), - pair_dim: wp.array(dtype=int), - pair_friction: wp.array2d(dtype=mjwp_types.vec5), - pair_gap: wp.array2d(dtype=float), - pair_margin: wp.array2d(dtype=float), - pair_solimp: wp.array2d(dtype=mjwp_types.vec5), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr), - qLD_all_updates: wp.array(dtype=wp.vec3i), - qLD_level_offsets: wp.array(dtype=int), - qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...], - qM_fullm_i: wp.array(dtype=int), - qM_fullm_j: wp.array(dtype=int), - qM_mulm_col: wp.array(dtype=int), - qM_mulm_madr: wp.array(dtype=int), - qM_mulm_rowadr: wp.array(dtype=int), + nxn_geom_pair_filtered: wp.array[wp.vec2i], + nxn_pairid: wp.array[wp.vec2i], + nxn_pairid_filtered: wp.array[wp.vec2i], + oct_aabb: wp.array2d[wp.vec3], + oct_child: wp.array[mjwp_types.vec8i], + oct_coeff: wp.array[mjwp_types.vec8], + pair_dim: wp.array[int], + pair_friction: wp.array2d[mjwp_types.vec5], + pair_gap: wp.array2d[float], + pair_margin: wp.array2d[float], + pair_solimp: wp.array2d[mjwp_types.vec5], + pair_solref: wp.array2d[wp.vec2], + pair_solreffriction: wp.array2d[wp.vec2], + plugin: wp.array[int], + plugin_attr: wp.array[mjwp_types.vec_pluginattr], + qLD_all_updates: wp.array[wp.vec3i], + qLD_level_offsets: wp.array[int], + qLD_updates: tuple[wp.array[wp.vec3i], ...], + qM_fullm_i: wp.array[int], + qM_fullm_j: wp.array[int], + qM_mulm_col: wp.array[int], + qM_mulm_madr: wp.array[int], + qM_mulm_rowadr: wp.array[int], qM_tiles: tuple[mjwp_types.TileSet, ...], - qpos0: wp.array2d(dtype=float), - qpos_spring: wp.array2d(dtype=float), - rangefinder_sensor_adr: wp.array(dtype=int), - sensor_acc_adr: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - sensor_adr_to_contact_adr: wp.array(dtype=int), - sensor_contact_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), - sensor_datatype: wp.array(dtype=int), - sensor_dim: wp.array(dtype=int), + qpos0: wp.array2d[float], + qpos_spring: wp.array2d[float], + rangefinder_sensor_adr: wp.array[int], + sensor_acc_adr: wp.array[int], + sensor_adr: wp.array[int], + sensor_adr_to_contact_adr: wp.array[int], + sensor_contact_adr: wp.array[int], + sensor_cutoff: wp.array[float], + sensor_datatype: wp.array[int], + sensor_dim: wp.array[int], sensor_e_kinetic: bool, sensor_e_potential: bool, - sensor_intprm: wp.array2d(dtype=int), - sensor_limitfrc_adr: wp.array(dtype=int), - sensor_limitpos_adr: wp.array(dtype=int), - sensor_limitvel_adr: wp.array(dtype=int), - sensor_objid: wp.array(dtype=int), - sensor_objtype: wp.array(dtype=int), - sensor_pos_adr: wp.array(dtype=int), - sensor_rangefinder_adr: wp.array(dtype=int), - sensor_rangefinder_bodyid: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_reftype: wp.array(dtype=int), + sensor_intprm: wp.array2d[int], + sensor_limitfrc_adr: wp.array[int], + sensor_limitpos_adr: wp.array[int], + sensor_limitvel_adr: wp.array[int], + sensor_objid: wp.array[int], + sensor_objtype: wp.array[int], + sensor_pos_adr: wp.array[int], + sensor_rangefinder_adr: wp.array[int], + sensor_rangefinder_bodyid: wp.array[int], + sensor_refid: wp.array[int], + sensor_reftype: wp.array[int], sensor_rne_postconstraint: bool, sensor_subtree_vel: bool, - sensor_tendonactfrc_adr: wp.array(dtype=int), - sensor_touch_adr: wp.array(dtype=int), - sensor_type: wp.array(dtype=int), - sensor_vel_adr: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_pos: wp.array2d(dtype=wp.vec3), - site_quat: wp.array2d(dtype=wp.quat), - site_size: wp.array(dtype=wp.vec3), - site_type: wp.array(dtype=int), - taxel_sensorid: wp.array(dtype=int), - taxel_vertadr: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - tendon_actfrclimited: wp.array(dtype=bool), - tendon_actfrcrange: wp.array2d(dtype=wp.vec2), - tendon_adr: wp.array(dtype=int), - tendon_armature: wp.array2d(dtype=float), - tendon_damping: wp.array2d(dtype=float), - tendon_frictionloss: wp.array2d(dtype=float), - tendon_geom_adr: wp.array(dtype=int), - tendon_invweight0: wp.array2d(dtype=float), - tendon_jnt_adr: wp.array(dtype=int), - tendon_length0: wp.array2d(dtype=float), - tendon_lengthspring: wp.array2d(dtype=wp.vec2), - tendon_limited_adr: wp.array(dtype=int), - tendon_margin: wp.array2d(dtype=float), - tendon_num: wp.array(dtype=int), - tendon_range: wp.array2d(dtype=wp.vec2), - tendon_site_pair_adr: wp.array(dtype=int), - tendon_solimp_fri: wp.array2d(dtype=mjwp_types.vec5), - tendon_solimp_lim: wp.array2d(dtype=mjwp_types.vec5), - tendon_solref_fri: wp.array2d(dtype=wp.vec2), - tendon_solref_lim: wp.array2d(dtype=wp.vec2), - tendon_stiffness: wp.array2d(dtype=float), - wrap_geom_adr: wp.array(dtype=int), - wrap_jnt_adr: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), - wrap_pulley_scale: wp.array(dtype=float), - wrap_site_pair_adr: wp.array(dtype=int), - wrap_type: wp.array(dtype=int), + sensor_tendonactfrc_adr: wp.array[int], + sensor_touch_adr: wp.array[int], + sensor_type: wp.array[int], + sensor_vel_adr: wp.array[int], + site_bodyid: wp.array[int], + site_pos: wp.array2d[wp.vec3], + site_quat: wp.array2d[wp.quat], + site_size: wp.array[wp.vec3], + site_type: wp.array[int], + taxel_sensorid: wp.array[int], + taxel_vertadr: wp.array[int], + ten_J_colind: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_rownnz: wp.array[int], + tendon_actfrclimited: wp.array[bool], + tendon_actfrcrange: wp.array2d[wp.vec2], + tendon_adr: wp.array[int], + tendon_armature: wp.array2d[float], + tendon_damping: wp.array2d[float], + tendon_frictionloss: wp.array2d[float], + tendon_geom_adr: wp.array[int], + tendon_invweight0: wp.array2d[float], + tendon_jnt_adr: wp.array[int], + tendon_length0: wp.array2d[float], + tendon_lengthspring: wp.array2d[wp.vec2], + tendon_limited_adr: wp.array[int], + tendon_margin: wp.array2d[float], + tendon_num: wp.array[int], + tendon_range: wp.array2d[wp.vec2], + tendon_site_pair_adr: wp.array[int], + tendon_solimp_fri: wp.array2d[mjwp_types.vec5], + tendon_solimp_lim: wp.array2d[mjwp_types.vec5], + tendon_solref_fri: wp.array2d[wp.vec2], + tendon_solref_lim: wp.array2d[wp.vec2], + tendon_stiffness: wp.array2d[float], + wrap_geom_adr: wp.array[int], + wrap_jnt_adr: wp.array[int], + wrap_objid: wp.array[int], + wrap_prm: wp.array[float], + wrap_pulley_scale: wp.array[float], + wrap_site_pair_adr: wp.array[int], + wrap_type: wp.array[int], opt__broadphase: int, opt__broadphase_filter: int, opt__ccd_iterations: int, - opt__ccd_tolerance: wp.array(dtype=float), + opt__ccd_tolerance: wp.array[float], opt__cone: int, opt__contact_sensor_maxmatch: int, - opt__density: wp.array(dtype=float), + opt__density: wp.array[float], opt__disableflags: int, opt__enableflags: int, opt__graph_conditional: bool, - opt__gravity: wp.array(dtype=wp.vec3), - opt__impratio_invsqrt: wp.array(dtype=float), + opt__gravity: wp.array[wp.vec3], + opt__impratio_invsqrt: wp.array[float], opt__integrator: int, opt__iterations: int, opt__ls_iterations: int, opt__ls_parallel: bool, opt__ls_parallel_min_step: float, - opt__ls_tolerance: wp.array(dtype=float), - opt__magnetic: wp.array(dtype=wp.vec3), + opt__ls_tolerance: wp.array[float], + opt__magnetic: wp.array[wp.vec3], opt__run_collision_detection: bool, opt__sdf_initpoints: int, opt__sdf_iterations: int, opt__solver: int, - opt__timestep: wp.array(dtype=float), - opt__tolerance: wp.array(dtype=float), - opt__viscosity: wp.array(dtype=float), - opt__wind: wp.array(dtype=wp.vec3), - stat__meaninertia: wp.array(dtype=float), + opt__timestep: wp.array[float], + opt__tolerance: wp.array[float], + opt__viscosity: wp.array[float], + opt__wind: wp.array[wp.vec3], + stat__meaninertia: wp.array[float], # Data naccdmax: int, naconmax: int, njmax: int, njmax_nnz: int, - act: wp.array2d(dtype=float), - act_dot: wp.array2d(dtype=float), - actuator_force: wp.array2d(dtype=float), - actuator_length: wp.array2d(dtype=float), - actuator_moment: wp.array2d(dtype=float), - actuator_velocity: wp.array2d(dtype=float), - cacc: wp.array2d(dtype=wp.spatial_vector), - cam_xmat: wp.array2d(dtype=wp.mat33), - cam_xpos: wp.array2d(dtype=wp.vec3), - cdof: wp.array2d(dtype=wp.spatial_vector), - cdof_dot: wp.array2d(dtype=wp.spatial_vector), - cfrc_ext: wp.array2d(dtype=wp.spatial_vector), - cfrc_int: wp.array2d(dtype=wp.spatial_vector), - cinert: wp.array2d(dtype=mjwp_types.vec10), - crb: wp.array2d(dtype=mjwp_types.vec10), - ctrl: wp.array2d(dtype=float), - cvel: wp.array2d(dtype=wp.spatial_vector), - energy: wp.array(dtype=wp.vec2), - eq_active: wp.array2d(dtype=bool), - flexedge_J: wp.array2d(dtype=float), - flexedge_length: wp.array2d(dtype=float), - flexedge_velocity: wp.array2d(dtype=float), - flexvert_xpos: wp.array2d(dtype=wp.vec3), - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), - light_xdir: wp.array2d(dtype=wp.vec3), - light_xpos: wp.array2d(dtype=wp.vec3), - mocap_pos: wp.array2d(dtype=wp.vec3), - mocap_quat: wp.array2d(dtype=wp.quat), - moment_colind: wp.array2d(dtype=int), - moment_rowadr: wp.array2d(dtype=int), - moment_rownnz: wp.array2d(dtype=int), - nacon: wp.array(dtype=int), - ncollision: wp.array(dtype=int), - ne: wp.array(dtype=int), - nefc: wp.array(dtype=int), - nf: wp.array(dtype=int), - nisland: wp.array(dtype=int), - nl: wp.array(dtype=int), - qLD: wp.array3d(dtype=float), - qLDiagInv: wp.array2d(dtype=float), - qM: wp.array3d(dtype=float), - qacc: wp.array2d(dtype=float), - qacc_smooth: wp.array2d(dtype=float), - qacc_warmstart: wp.array2d(dtype=float), - qfrc_actuator: wp.array2d(dtype=float), - qfrc_applied: wp.array2d(dtype=float), - qfrc_bias: wp.array2d(dtype=float), - qfrc_constraint: wp.array2d(dtype=float), - qfrc_damper: wp.array2d(dtype=float), - qfrc_fluid: wp.array2d(dtype=float), - qfrc_gravcomp: wp.array2d(dtype=float), - qfrc_passive: wp.array2d(dtype=float), - qfrc_smooth: wp.array2d(dtype=float), - qfrc_spring: wp.array2d(dtype=float), - qpos: wp.array2d(dtype=float), - qvel: wp.array2d(dtype=float), - sensordata: wp.array2d(dtype=float), - site_xmat: wp.array2d(dtype=wp.mat33), - site_xpos: wp.array2d(dtype=wp.vec3), - solver_niter: wp.array(dtype=int), - subtree_angmom: wp.array2d(dtype=wp.vec3), - subtree_com: wp.array2d(dtype=wp.vec3), - subtree_linvel: wp.array2d(dtype=wp.vec3), - ten_J: wp.array2d(dtype=float), - ten_length: wp.array2d(dtype=float), - ten_velocity: wp.array2d(dtype=float), - ten_wrapadr: wp.array2d(dtype=int), - ten_wrapnum: wp.array2d(dtype=int), - time: wp.array(dtype=float), - tree_island: wp.array2d(dtype=int), - wrap_obj: wp.array2d(dtype=wp.vec2i), - wrap_xpos: wp.array2d(dtype=wp.spatial_vector), - xanchor: wp.array2d(dtype=wp.vec3), - xaxis: wp.array2d(dtype=wp.vec3), - xfrc_applied: wp.array2d(dtype=wp.spatial_vector), - ximat: wp.array2d(dtype=wp.mat33), - xipos: wp.array2d(dtype=wp.vec3), - xmat: wp.array2d(dtype=wp.mat33), - xpos: wp.array2d(dtype=wp.vec3), - xquat: wp.array2d(dtype=wp.quat), - contact__dim: wp.array(dtype=int), - contact__dist: wp.array(dtype=float), - contact__efc_address: wp.array2d(dtype=int), - contact__flex: wp.array(dtype=wp.vec2i), - contact__frame: wp.array(dtype=wp.mat33), - contact__friction: wp.array(dtype=mjwp_types.vec5), - contact__geom: wp.array(dtype=wp.vec2i), - contact__geomcollisionid: wp.array(dtype=int), - contact__includemargin: wp.array(dtype=float), - contact__pos: wp.array(dtype=wp.vec3), - contact__solimp: wp.array(dtype=mjwp_types.vec5), - contact__solref: wp.array(dtype=wp.vec2), - contact__solreffriction: wp.array(dtype=wp.vec2), - contact__type: wp.array(dtype=int), - contact__vert: wp.array(dtype=wp.vec2i), - contact__worldid: wp.array(dtype=int), - efc__D: wp.array2d(dtype=float), - efc__J: wp.array3d(dtype=float), - efc__J_colind: wp.array3d(dtype=int), - efc__J_rowadr: wp.array2d(dtype=int), - efc__J_rownnz: wp.array2d(dtype=int), - efc__Ma: wp.array2d(dtype=float), - efc__aref: wp.array2d(dtype=float), - efc__force: wp.array2d(dtype=float), - efc__frictionloss: wp.array2d(dtype=float), - efc__id: wp.array2d(dtype=int), - efc__margin: wp.array2d(dtype=float), - efc__pos: wp.array2d(dtype=float), - efc__state: wp.array2d(dtype=int), - efc__type: wp.array2d(dtype=int), - efc__vel: wp.array2d(dtype=float), + act: wp.array2d[float], + act_dot: wp.array2d[float], + actuator_force: wp.array2d[float], + actuator_length: wp.array2d[float], + actuator_moment: wp.array2d[float], + actuator_velocity: wp.array2d[float], + cacc: wp.array2d[wp.spatial_vector], + cam_xmat: wp.array2d[wp.mat33], + cam_xpos: wp.array2d[wp.vec3], + cdof: wp.array2d[wp.spatial_vector], + cdof_dot: wp.array2d[wp.spatial_vector], + cfrc_ext: wp.array2d[wp.spatial_vector], + cfrc_int: wp.array2d[wp.spatial_vector], + cinert: wp.array2d[mjwp_types.vec10], + crb: wp.array2d[mjwp_types.vec10], + ctrl: wp.array2d[float], + cvel: wp.array2d[wp.spatial_vector], + energy: wp.array[wp.vec2], + eq_active: wp.array2d[bool], + flexedge_J: wp.array2d[float], + flexedge_length: wp.array2d[float], + flexedge_velocity: wp.array2d[float], + flexvert_xpos: wp.array2d[wp.vec3], + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], + light_xdir: wp.array2d[wp.vec3], + light_xpos: wp.array2d[wp.vec3], + mocap_pos: wp.array2d[wp.vec3], + mocap_quat: wp.array2d[wp.quat], + moment_colind: wp.array2d[int], + moment_rowadr: wp.array2d[int], + moment_rownnz: wp.array2d[int], + nacon: wp.array[int], + ncollision: wp.array[int], + ne: wp.array[int], + nefc: wp.array[int], + nf: wp.array[int], + nisland: wp.array[int], + nl: wp.array[int], + qLD: wp.array3d[float], + qLDiagInv: wp.array2d[float], + qM: wp.array3d[float], + qacc: wp.array2d[float], + qacc_smooth: wp.array2d[float], + qacc_warmstart: wp.array2d[float], + qfrc_actuator: wp.array2d[float], + qfrc_applied: wp.array2d[float], + qfrc_bias: wp.array2d[float], + qfrc_constraint: wp.array2d[float], + qfrc_damper: wp.array2d[float], + qfrc_fluid: wp.array2d[float], + qfrc_gravcomp: wp.array2d[float], + qfrc_passive: wp.array2d[float], + qfrc_smooth: wp.array2d[float], + qfrc_spring: wp.array2d[float], + qpos: wp.array2d[float], + qvel: wp.array2d[float], + sensordata: wp.array2d[float], + site_xmat: wp.array2d[wp.mat33], + site_xpos: wp.array2d[wp.vec3], + solver_niter: wp.array[int], + subtree_angmom: wp.array2d[wp.vec3], + subtree_com: wp.array2d[wp.vec3], + subtree_linvel: wp.array2d[wp.vec3], + ten_J: wp.array2d[float], + ten_length: wp.array2d[float], + ten_velocity: wp.array2d[float], + ten_wrapadr: wp.array2d[int], + ten_wrapnum: wp.array2d[int], + time: wp.array[float], + tree_island: wp.array2d[int], + wrap_obj: wp.array2d[wp.vec2i], + wrap_xpos: wp.array2d[wp.spatial_vector], + xanchor: wp.array2d[wp.vec3], + xaxis: wp.array2d[wp.vec3], + xfrc_applied: wp.array2d[wp.spatial_vector], + ximat: wp.array2d[wp.mat33], + xipos: wp.array2d[wp.vec3], + xmat: wp.array2d[wp.mat33], + xpos: wp.array2d[wp.vec3], + xquat: wp.array2d[wp.quat], + contact__dim: wp.array[int], + contact__dist: wp.array[float], + contact__efc_address: wp.array2d[int], + contact__flex: wp.array[wp.vec2i], + contact__frame: wp.array[wp.mat33], + contact__friction: wp.array[mjwp_types.vec5], + contact__geom: wp.array[wp.vec2i], + contact__geomcollisionid: wp.array[int], + contact__includemargin: wp.array[float], + contact__pos: wp.array[wp.vec3], + contact__solimp: wp.array[mjwp_types.vec5], + contact__solref: wp.array[wp.vec2], + contact__solreffriction: wp.array[wp.vec2], + contact__type: wp.array[int], + contact__vert: wp.array[wp.vec2i], + contact__worldid: wp.array[int], + efc__D: wp.array2d[float], + efc__J: wp.array3d[float], + efc__J_colind: wp.array3d[int], + efc__J_rowadr: wp.array2d[int], + efc__J_rownnz: wp.array2d[int], + efc__Ma: wp.array2d[float], + efc__aref: wp.array2d[float], + efc__force: wp.array2d[float], + efc__frictionloss: wp.array2d[float], + efc__id: wp.array2d[int], + efc__margin: wp.array2d[float], + efc__pos: wp.array2d[float], + efc__state: wp.array2d[int], + efc__type: wp.array2d[int], + efc__vel: wp.array2d[float], ): _m.stat = _s _m.opt = _o @@ -3457,7 +3458,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.geom_conaffinity, m.geom_condim, m.geom_contype, - m.geom_dataid, + jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_fluid, m.geom_friction, m.geom_gap, diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 6a3f66ab..c97003f0 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -48,42 +48,43 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) + @ffi.format_args_for_warp def _render_shim( # Model nworld: int, - cam_fovy: wp.array2d(dtype=float), - cam_intrinsic: wp.array2d(dtype=wp.vec4), - cam_projection: wp.array(dtype=int), - cam_sensorsize: wp.array(dtype=wp.vec2), - flex_edge: wp.array(dtype=wp.vec2i), - flex_radius: wp.array(dtype=float), - flex_vertadr: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_matid: wp.array2d(dtype=int), - geom_rgba: wp.array2d(dtype=wp.vec4), - geom_size: wp.array2d(dtype=wp.vec3), - geom_type: wp.array(dtype=int), - light_active: wp.array2d(dtype=bool), - light_castshadow: wp.array2d(dtype=bool), - light_type: wp.array2d(dtype=int), - mat_rgba: wp.array2d(dtype=wp.vec4), - mat_texid: wp.array3d(dtype=int), - mat_texrepeat: wp.array2d(dtype=wp.vec2), - mesh_faceadr: wp.array(dtype=int), + cam_fovy: wp.array2d[float], + cam_intrinsic: wp.array2d[wp.vec4], + cam_projection: wp.array[int], + cam_sensorsize: wp.array[wp.vec2], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], + flex_vertadr: wp.array[int], + geom_dataid: wp.array2d[int], + geom_matid: wp.array2d[int], + geom_rgba: wp.array2d[wp.vec4], + geom_size: wp.array2d[wp.vec3], + geom_type: wp.array[int], + light_active: wp.array2d[bool], + light_castshadow: wp.array2d[bool], + light_type: wp.array2d[int], + mat_rgba: wp.array2d[wp.vec4], + mat_texid: wp.array3d[int], + mat_texrepeat: wp.array2d[wp.vec2], + mesh_faceadr: wp.array[int], nlight: int, # Data - cam_xmat: wp.array2d(dtype=wp.mat33), - cam_xpos: wp.array2d(dtype=wp.vec3), - flexvert_xpos: wp.array2d(dtype=wp.vec3), - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), - light_xdir: wp.array2d(dtype=wp.vec3), - light_xpos: wp.array2d(dtype=wp.vec3), + cam_xmat: wp.array2d[wp.mat33], + cam_xpos: wp.array2d[wp.vec3], + flexvert_xpos: wp.array2d[wp.vec3], + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], + light_xdir: wp.array2d[wp.vec3], + light_xpos: wp.array2d[wp.vec3], # Registry rc_id: int, - rgb: wp.array2d(dtype=wp.uint32), - depth: wp.array2d(dtype=wp.float32), + rgb: wp.array2d[wp.uint32], + depth: wp.array2d[wp.float32], ): _m.stat = _s _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_radius, m.flex_vertadr, - m.geom_dataid, + jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_matid, m.geom_rgba, m.geom_size, diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 209dc6b0..020b9392 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -46,52 +46,53 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) + @ffi.format_args_for_warp def _kinematics_shim( # Model nworld: int, - body_branch_start: wp.array(dtype=int), - body_branches: wp.array(dtype=int), - body_ipos: wp.array2d(dtype=wp.vec3), - body_iquat: wp.array2d(dtype=wp.quat), - body_jntadr: wp.array(dtype=int), - body_jntnum: wp.array(dtype=int), - body_mocapid: wp.array(dtype=int), - body_parentid: wp.array(dtype=int), - body_pos: wp.array2d(dtype=wp.vec3), - body_quat: wp.array2d(dtype=wp.quat), - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_pos: wp.array2d(dtype=wp.vec3), - geom_quat: wp.array2d(dtype=wp.quat), - jnt_axis: wp.array2d(dtype=wp.vec3), - jnt_pos: wp.array2d(dtype=wp.vec3), - jnt_qposadr: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), + body_branch_start: wp.array[int], + body_branches: wp.array[int], + body_ipos: wp.array2d[wp.vec3], + body_iquat: wp.array2d[wp.quat], + body_jntadr: wp.array[int], + body_jntnum: wp.array[int], + body_mocapid: wp.array[int], + body_parentid: wp.array[int], + body_pos: wp.array2d[wp.vec3], + body_quat: wp.array2d[wp.quat], + body_rootid: wp.array[int], + body_weldid: wp.array[int], + geom_bodyid: wp.array[int], + geom_pos: wp.array2d[wp.vec3], + geom_quat: wp.array2d[wp.quat], + jnt_axis: wp.array2d[wp.vec3], + jnt_pos: wp.array2d[wp.vec3], + jnt_qposadr: wp.array[int], + jnt_type: wp.array[int], nbody: int, nbranch: int, ngeom: int, nsite: int, - qpos0: wp.array2d(dtype=float), - site_bodyid: wp.array(dtype=int), - site_pos: wp.array2d(dtype=wp.vec3), - site_quat: wp.array2d(dtype=wp.quat), + qpos0: wp.array2d[float], + site_bodyid: wp.array[int], + site_pos: wp.array2d[wp.vec3], + site_quat: wp.array2d[wp.quat], # Data - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), - mocap_pos: wp.array2d(dtype=wp.vec3), - mocap_quat: wp.array2d(dtype=wp.quat), - qpos: wp.array2d(dtype=float), - site_xmat: wp.array2d(dtype=wp.mat33), - site_xpos: wp.array2d(dtype=wp.vec3), - xanchor: wp.array2d(dtype=wp.vec3), - xaxis: wp.array2d(dtype=wp.vec3), - ximat: wp.array2d(dtype=wp.mat33), - xipos: wp.array2d(dtype=wp.vec3), - xmat: wp.array2d(dtype=wp.mat33), - xpos: wp.array2d(dtype=wp.vec3), - xquat: wp.array2d(dtype=wp.quat), + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], + mocap_pos: wp.array2d[wp.vec3], + mocap_quat: wp.array2d[wp.quat], + qpos: wp.array2d[float], + site_xmat: wp.array2d[wp.mat33], + site_xpos: wp.array2d[wp.vec3], + xanchor: wp.array2d[wp.vec3], + xaxis: wp.array2d[wp.vec3], + ximat: wp.array2d[wp.mat33], + xipos: wp.array2d[wp.vec3], + xmat: wp.array2d[wp.mat33], + xpos: wp.array2d[wp.vec3], + xquat: wp.array2d[wp.quat], ): _m.stat = _s _m.opt = _o @@ -297,45 +298,45 @@ def kinematics_vmap( def _tendon_shim( # Model nworld: int, - body_dofadr: wp.array(dtype=int), - body_dofnum: wp.array(dtype=int), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - jnt_dofadr: wp.array(dtype=int), - jnt_qposadr: wp.array(dtype=int), + body_dofadr: wp.array[int], + body_dofnum: wp.array[int], + body_parentid: wp.array[int], + body_rootid: wp.array[int], + geom_bodyid: wp.array[int], + geom_size: wp.array2d[wp.vec3], + jnt_dofadr: wp.array[int], + jnt_qposadr: wp.array[int], ntendon: int, nwrap: int, - site_bodyid: wp.array(dtype=int), - ten_J_colind: wp.array(dtype=int), - ten_J_rowadr: wp.array(dtype=int), - ten_J_rownnz: wp.array(dtype=int), - tendon_adr: wp.array(dtype=int), - tendon_geom_adr: wp.array(dtype=int), - tendon_jnt_adr: wp.array(dtype=int), - tendon_num: wp.array(dtype=int), - tendon_site_pair_adr: wp.array(dtype=int), - wrap_geom_adr: wp.array(dtype=int), - wrap_jnt_adr: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), - wrap_pulley_scale: wp.array(dtype=float), - wrap_site_pair_adr: wp.array(dtype=int), - wrap_type: wp.array(dtype=int), + site_bodyid: wp.array[int], + ten_J_colind: wp.array[int], + ten_J_rowadr: wp.array[int], + ten_J_rownnz: wp.array[int], + tendon_adr: wp.array[int], + tendon_geom_adr: wp.array[int], + tendon_jnt_adr: wp.array[int], + tendon_num: wp.array[int], + tendon_site_pair_adr: wp.array[int], + wrap_geom_adr: wp.array[int], + wrap_jnt_adr: wp.array[int], + wrap_objid: wp.array[int], + wrap_prm: wp.array[float], + wrap_pulley_scale: wp.array[float], + wrap_site_pair_adr: wp.array[int], + wrap_type: wp.array[int], # Data - cdof: wp.array2d(dtype=wp.spatial_vector), - geom_xmat: wp.array2d(dtype=wp.mat33), - geom_xpos: wp.array2d(dtype=wp.vec3), - qpos: wp.array2d(dtype=float), - site_xpos: wp.array2d(dtype=wp.vec3), - subtree_com: wp.array2d(dtype=wp.vec3), - ten_J: wp.array2d(dtype=float), - ten_length: wp.array2d(dtype=float), - ten_wrapadr: wp.array2d(dtype=int), - ten_wrapnum: wp.array2d(dtype=int), - wrap_obj: wp.array2d(dtype=wp.vec2i), - wrap_xpos: wp.array2d(dtype=wp.spatial_vector), + cdof: wp.array2d[wp.spatial_vector], + geom_xmat: wp.array2d[wp.mat33], + geom_xpos: wp.array2d[wp.vec3], + qpos: wp.array2d[float], + site_xpos: wp.array2d[wp.vec3], + subtree_com: wp.array2d[wp.vec3], + ten_J: wp.array2d[float], + ten_length: wp.array2d[float], + ten_wrapadr: wp.array2d[int], + ten_wrapnum: wp.array2d[int], + wrap_obj: wp.array2d[wp.vec2i], + wrap_xpos: wp.array2d[wp.spatial_vector], ): _m.stat = _s _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( # Model nworld: int, - body_inertia: wp.array2d(dtype=wp.vec3), - body_mass: wp.array2d(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_subtreemass: wp.array2d(dtype=float), - body_tree: tuple[wp.array(dtype=int), ...], - jnt_bodyid: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_type: wp.array(dtype=int), + body_inertia: wp.array2d[wp.vec3], + body_mass: wp.array2d[float], + body_parentid: wp.array[int], + body_rootid: wp.array[int], + body_subtreemass: wp.array2d[float], + body_tree: tuple[wp.array[int], ...], + jnt_bodyid: wp.array[int], + jnt_dofadr: wp.array[int], + jnt_type: wp.array[int], nbody: int, njnt: int, # Data - cdof: wp.array2d(dtype=wp.spatial_vector), - cinert: wp.array2d(dtype=mjwp_types.vec10), - subtree_com: wp.array2d(dtype=wp.vec3), - xanchor: wp.array2d(dtype=wp.vec3), - xaxis: wp.array2d(dtype=wp.vec3), - ximat: wp.array2d(dtype=wp.mat33), - xipos: wp.array2d(dtype=wp.vec3), - xmat: wp.array2d(dtype=wp.mat33), + cdof: wp.array2d[wp.spatial_vector], + cinert: wp.array2d[mjwp_types.vec10], + subtree_com: wp.array2d[wp.vec3], + xanchor: wp.array2d[wp.vec3], + xaxis: wp.array2d[wp.vec3], + ximat: wp.array2d[wp.mat33], + xipos: wp.array2d[wp.vec3], + xmat: wp.array2d[wp.mat33], ): _m.stat = _s _m.opt = _o diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 564fe57e..14e416a8 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -672,7 +672,7 @@ _NDIM = { 'geom_conaffinity': 1, 'geom_condim': 1, 'geom_contype': 1, - 'geom_dataid': 1, + 'geom_dataid': 2, 'geom_fluid': 2, 'geom_friction': 3, 'geom_gap': 2, @@ -1238,7 +1238,7 @@ _BATCH_DIM = { 'geom_conaffinity': False, 'geom_condim': False, 'geom_contype': False, - 'geom_dataid': False, + 'geom_dataid': True, 'geom_fluid': False, 'geom_friction': True, 'geom_gap': True,