diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 54e40926..d9ef50b3 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -205,6 +205,9 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any: ) if isinstance(wp_field, mjwp_types.BlockDim): return mjxw.types.BlockDim(**wp_field.__dict__) + if isinstance(wp_field, mjwp_types.Callback): + return wp_field + if isinstance(wp_field, tuple) and is_static(wp_field[0]): return wp_field diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 5b8ad301..f4cda34c 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -848,9 +848,9 @@ class Model(PyTreeNode): actuator_forcerange: jax.Array actuator_actrange: jax.Array actuator_gear: jax.Array - actuator_cranklength: np.ndarray + actuator_cranklength: jax.Array actuator_acc0: jax.Array - actuator_lengthrange: np.ndarray + actuator_lengthrange: jax.Array sensor_type: np.ndarray sensor_datatype: np.ndarray sensor_needstage: np.ndarray diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py index 936bd9fa..1ff05ff6 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py @@ -56,6 +56,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.io import reset_data as reset_data from mujoco.mjx.third_party.mujoco_warp._src.io import set_const as set_const from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_0 as set_const_0 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_fixed as set_const_fixed +from mujoco.mjx.third_party.mujoco_warp._src.io import set_length_range as set_length_range +from mujoco.mjx.third_party.mujoco_warp._src.island import island as island from mujoco.mjx.third_party.mujoco_warp._src.passive import passive as passive from mujoco.mjx.third_party.mujoco_warp._src.ray import ray as ray from mujoco.mjx.third_party.mujoco_warp._src.ray import rays as rays 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 42912470..5712b81f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py @@ -57,12 +57,12 @@ def ctrl_noise( worldid, actid = wp.tid() # convert rate and scale to discrete time (Ornstein-Uhlenbeck) - rate = wp.exp(-opt_timestep[0] / ctrlnoiserate) + rate = wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / ctrlnoiserate) scale = ctrlnoisestd * wp.sqrt(1.0 - rate * rate) midpoint = 0.0 halfrange = 1.0 - ctrlrange = actuator_ctrlrange[0, actid] + ctrlrange = actuator_ctrlrange[worldid % actuator_ctrlrange.shape[0], actid] is_limited = actuator_ctrllimited[actid] if is_limited: midpoint = 0.5 * (ctrlrange[1] + ctrlrange[0]) 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 7cddcacc..c40fcfc9 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py @@ -215,7 +215,8 @@ def _compute_bvh_bounds( lower_bound, upper_bound = _compute_box_bounds(pos, rot, size) elif type == GeomType.HFIELD: size = hfield_bounds_size[geom_dataid[geom_id]] - lower_bound, upper_bound = _compute_box_bounds(pos, rot, size) + hfield_center = pos + rot[:, 2] * size[2] + lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size) lower_out[world_id * bvh_ngeom + geom_local_id] = lower_bound upper_out[world_id * bvh_ngeom + geom_local_id] = upper_bound 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 a8bd2d8a..69ea9cbe 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 @@ -15,32 +15,34 @@ from typing import Tuple -import warp as wp - +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support -from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params +from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact -from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index +from mujoco.mjx.third_party.mujoco_warp._src.types import Data +from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit +from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType +from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 +from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL -from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext -from mujoco.mjx.third_party.mujoco_warp._src.types import Data -from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit -from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import Model -from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 -from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) @@ -92,10 +94,7 @@ def _hfield_filter( r2 = geom_rbound[rbound_id, g2] # TODO(team): margin? - if BLEEDING_EDGE_MUJOCO: - margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2] - else: - margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2]) + margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2] # box-sphere test: horizontal plane for i in range(2): @@ -789,6 +788,13 @@ def ccd_kernel_builder( if dist >= 0.0 and pairid[1] == -1: return 0 + # CCD operates on margin-inflated shapes (support() inflates each geom by + # 0.5 * margin). The returned dist is therefore relative to the inflated + # geometry. Correct back to the true surface-to-surface distance so that + # the constraint pipeline (pos = dist - includemargin) works consistently + # with the primitive narrowphase, which reports un-inflated distances. + dist += margin + witness1[0] = w1 witness2[0] = w2 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 new file mode 100644 index 00000000..f8294f65 --- /dev/null +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py @@ -0,0 +1,372 @@ +# Copyright 2026 The Newton Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Core collision types and utilities shared across collision modules.""" + +import dataclasses +from typing import Tuple + +from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div +from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType +from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType +from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL +from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 +import warp as wp + +wp.set_module_options({"enable_backward": False}) + + +@wp.struct +class Geom: + """Geom properties for pairwise collision detection. + + Bundles a geometry's pose, size, surface normal, and mesh topology data into + a single struct that can be passed to Warp collision kernels. + """ + + pos: wp.vec3 + rot: wp.mat33 + normal: wp.vec3 + size: wp.vec3 + margin: float + hfprism: mat63 + vertadr: int + vertnum: int + vert: wp.array(dtype=wp.vec3) + graphadr: int + graph: wp.array(dtype=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) + 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), + # Data in: + geom_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xmat_in: wp.array2d(dtype=wp.mat33), + # In: + geoms: wp.vec2i, + worldid: int, +) -> Tuple[Geom, Geom]: + geom1 = Geom() + geom2 = Geom() + + g1 = geoms[0] + g2 = geoms[1] + geom_type1 = geom_type[g1] + geom_type2 = geom_type[g2] + + geom1.pos = geom_xpos_in[worldid, g1] + geom1.rot = geom_xmat_in[worldid, g1] + geom1.size = geom_size[worldid % geom_size.shape[0], g1] + # z-axis of the rotation matrix, used as the surface normal for plane collisions + geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2]) + + geom2.pos = geom_xpos_in[worldid, g2] + geom2.rot = geom_xmat_in[worldid, g2] + geom2.size = geom_size[worldid % geom_size.shape[0], g2] + # 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]) + + if geom_type1 == GeomType.MESH: + dataid = geom_dataid[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) + geom1.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1) + geom1.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1) + + geom1.vert = mesh_vert + geom1.graph = mesh_graph + geom1.mesh_polynormal = mesh_polynormal + geom1.mesh_polyvertadr = mesh_polyvertadr + geom1.mesh_polyvertnum = mesh_polyvertnum + geom1.mesh_polyvert = mesh_polyvert + geom1.mesh_polymapadr = mesh_polymapadr + geom1.mesh_polymapnum = mesh_polymapnum + geom1.mesh_polymap = mesh_polymap + + if geom_type2 == GeomType.MESH: + dataid = geom_dataid[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) + geom2.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1) + geom2.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1) + + geom2.vert = mesh_vert + geom2.graph = mesh_graph + geom2.mesh_polynormal = mesh_polynormal + geom2.mesh_polyvertadr = mesh_polyvertadr + geom2.mesh_polyvertnum = mesh_polyvertnum + geom2.mesh_polyvert = mesh_polyvert + geom2.mesh_polymapadr = mesh_polymapadr + geom2.mesh_polymapnum = mesh_polymapnum + geom2.mesh_polymap = mesh_polymap + + geom1.index = -1 + geom1.margin = 0.0 + + geom2.index = -1 + geom2.margin = 0.0 + + return geom1, geom2 + + +@wp.func +def write_contact( + # Data in: + naconmax_in: int, + # In: + id_: int, + dist_in: float, + pos_in: wp.vec3, + frame_in: wp.mat33, + margin_in: float, + gap_in: float, + condim_in: int, + friction_in: vec5, + solref_in: wp.vec2, + solreffriction_in: wp.vec2, + solimp_in: vec5, + geoms_in: wp.vec2i, + pairid_in: wp.vec2i, + worldid_in: int, + # Data out: + contact_dist_out: wp.array(dtype=float), + contact_pos_out: wp.array(dtype=wp.vec3), + contact_frame_out: wp.array(dtype=wp.mat33), + contact_includemargin_out: wp.array(dtype=float), + contact_friction_out: wp.array(dtype=vec5), + contact_solref_out: wp.array(dtype=wp.vec2), + contact_solreffriction_out: wp.array(dtype=wp.vec2), + contact_solimp_out: wp.array(dtype=vec5), + contact_dim_out: wp.array(dtype=int), + contact_geom_out: wp.array(dtype=wp.vec2i), + contact_worldid_out: wp.array(dtype=int), + contact_type_out: wp.array(dtype=int), + contact_geomcollisionid_out: wp.array(dtype=int), + nacon_out: wp.array(dtype=int), +) -> int: + """Atomically write a detected contact into the contact output arrays. + + Returns 1 if the contact is active (dist < margin), 0 otherwise. + """ + active = dist_in < margin_in + + # skip contact and no collision sensor + if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1: + return 0 + + contact_type = 0 + + if pairid_in[0] >= -1 and active: + contact_type |= ContactType.CONSTRAINT + + if pairid_in[1] >= 0: + contact_type |= ContactType.SENSOR + + cid = wp.atomic_add(nacon_out, 0, 1) + if cid < naconmax_in: + contact_dist_out[cid] = dist_in + contact_pos_out[cid] = pos_in + contact_frame_out[cid] = frame_in + contact_geom_out[cid] = geoms_in + contact_worldid_out[cid] = worldid_in + includemargin = margin_in - gap_in + contact_includemargin_out[cid] = includemargin + contact_dim_out[cid] = condim_in + contact_friction_out[cid] = friction_in + contact_solref_out[cid] = solref_in + contact_solreffriction_out[cid] = solreffriction_in + contact_solimp_out[cid] = solimp_in + contact_type_out[cid] = contact_type + contact_geomcollisionid_out[cid] = id_ + return int(active) + return 0 + + +@wp.func +def contact_params( + # Model: + geom_condim: wp.array(dtype=int), + geom_priority: wp.array(dtype=int), + geom_solmix: wp.array2d(dtype=float), + geom_solref: wp.array2d(dtype=wp.vec2), + geom_solimp: wp.array2d(dtype=vec5), + geom_friction: wp.array2d(dtype=wp.vec3), + geom_margin: wp.array2d(dtype=float), + geom_gap: wp.array2d(dtype=float), + pair_dim: wp.array(dtype=int), + pair_solref: wp.array2d(dtype=wp.vec2), + pair_solreffriction: wp.array2d(dtype=wp.vec2), + pair_solimp: wp.array2d(dtype=vec5), + pair_margin: wp.array2d(dtype=float), + pair_gap: wp.array2d(dtype=float), + pair_friction: wp.array2d(dtype=vec5), + # In: + collision_pair_in: wp.array(dtype=wp.vec2i), + collision_pairid_in: wp.array(dtype=wp.vec2i), + cid: int, + worldid: int, +): + """Resolve contact parameters for a collision pair. + + Uses explicit pair overrides when available, otherwise mixes geom-level + properties by priority and solmix weights. + """ + geoms = collision_pair_in[cid] + pairid = collision_pairid_in[cid][0] + + # TODO(team): early return if collision sensor but no contact + # (ie, pairid[0] < -1 and pairid[1] < 0) + + if pairid > -1: + margin = pair_margin[worldid % pair_margin.shape[0], pairid] + gap = pair_gap[worldid % pair_gap.shape[0], pairid] + condim = pair_dim[pairid] + friction = pair_friction[worldid % pair_friction.shape[0], pairid] + solref = pair_solref[worldid % pair_solref.shape[0], pairid] + solreffriction = pair_solreffriction[ + worldid % pair_solreffriction.shape[0], pairid + ] + solimp = pair_solimp[worldid % pair_solimp.shape[0], pairid] + else: + g1 = geoms[0] + g2 = geoms[1] + solmix_id = worldid % geom_solmix.shape[0] + friction_id = worldid % geom_friction.shape[0] + solref_id = worldid % geom_solref.shape[0] + solimp_id = worldid % geom_solimp.shape[0] + margin_id = worldid % geom_margin.shape[0] + gap_id = worldid % geom_gap.shape[0] + + solmix1 = geom_solmix[solmix_id, g1] + solmix2 = geom_solmix[solmix_id, g2] + + condim1 = geom_condim[g1] + condim2 = geom_condim[g2] + + # priority + p1 = geom_priority[g1] + p2 = geom_priority[g2] + + if p1 > p2: + mix = 1.0 + condim = condim1 + max_geom_friction = geom_friction[friction_id, g1] + elif p2 > p1: + mix = 0.0 + condim = condim2 + max_geom_friction = geom_friction[friction_id, g2] + else: + mix = safe_div(solmix1, solmix1 + solmix2) + mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix) + mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix) + mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix) + condim = wp.max(condim1, condim2) + max_geom_friction = wp.max( + geom_friction[friction_id, g1], geom_friction[friction_id, g2] + ) + + friction = vec5( + max_geom_friction[0], + max_geom_friction[0], + max_geom_friction[1], + max_geom_friction[2], + max_geom_friction[2], + ) + + if ( + geom_solref[solref_id, g1][0] > 0.0 + and geom_solref[solref_id, g2][0] > 0.0 + ): + solref = ( + mix * geom_solref[solref_id, g1] + + (1.0 - mix) * geom_solref[solref_id, g2] + ) + else: + solref = wp.min(geom_solref[solref_id, g1], geom_solref[solref_id, g2]) + + solreffriction = wp.vec2(0.0, 0.0) + solimp = ( + mix * geom_solimp[solimp_id, g1] + + (1.0 - mix) * geom_solimp[solimp_id, g2] + ) + # geom priority is ignored + margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2] + gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2] + + friction = vec5( + wp.max(MJ_MINMU, friction[0]), + wp.max(MJ_MINMU, friction[1]), + wp.max(MJ_MINMU, friction[2]), + wp.max(MJ_MINMU, friction[3]), + wp.max(MJ_MINMU, friction[4]), + ) + + return geoms, margin, gap, condim, friction, solref, solreffriction, solimp + + +@dataclasses.dataclass +class CollisionContext: + """Collision driver intermediate arrays. + + Attributes: + collision_pair: collision pairs from broadphase (naconmax, 2) + collision_pairid: ids from broadphase (naconmax, 2) + collision_worldid: collision world ids from broadphase (naconmax,) + """ + + collision_pair: wp.array + collision_pairid: wp.array + collision_worldid: wp.array + + +def create_collision_context(naconmax: int) -> CollisionContext: + """Create a CollisionContext with allocated arrays.""" + return CollisionContext( + collision_pair=wp.empty(naconmax, dtype=wp.vec2i), + collision_pairid=wp.empty(naconmax, dtype=wp.vec2i), + collision_worldid=wp.empty(naconmax, dtype=int), + ) 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 0cb971dc..36cdaea1 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 @@ -15,26 +15,25 @@ from typing import Any -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src.collision_convex import convex_narrowphase +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import create_collision_context from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import primitive_narrowphase from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf_narrowphase -from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO from mujoco.mjx.third_party.mujoco_warp._src.math import upper_tri_index -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType -from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType -from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import mat23 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL +from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -76,15 +75,6 @@ MJ_COLLISION_TABLE = { } -def create_collision_context(naconmax: int) -> CollisionContext: - """Create a CollisionContext with allocated arrays.""" - return CollisionContext( - collision_pair=wp.empty(naconmax, dtype=wp.vec2i), - collision_pairid=wp.empty(naconmax, dtype=wp.vec2i), - collision_worldid=wp.empty(naconmax, dtype=int), - ) - - @wp.kernel def _zero_nacon_ncollision( # Data out: @@ -102,27 +92,18 @@ def _plane_filter( if size1 == 0.0: # geom1 is a plane dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2])) - if BLEEDING_EDGE_MUJOCO: - return dist <= size2 + margin1 + margin2 - else: - return dist <= size2 + wp.max(margin1, margin2) + return dist <= size2 + margin1 + margin2 elif size2 == 0.0: # geom2 is a plane dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2])) - if BLEEDING_EDGE_MUJOCO: - return dist <= size1 + margin1 + margin2 - else: - return dist <= size1 + wp.max(margin1, margin2) + return dist <= size1 + margin1 + margin2 return True @wp.func def _sphere_filter(size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3) -> bool: - if BLEEDING_EDGE_MUJOCO: - bound = size1 + size2 + margin1 + margin2 - else: - bound = size1 + size2 + wp.max(margin1, margin2) + bound = size1 + size2 + margin1 + margin2 dif = xpos2 - xpos1 dist_sq = wp.dot(dif, dif) return dist_sq <= bound * bound @@ -151,10 +132,7 @@ def _aabb_filter( center1 = xmat1 @ center1 + xpos1 center2 = xmat2 @ center2 + xpos2 - if BLEEDING_EDGE_MUJOCO: - margin = margin1 + margin2 - else: - margin = wp.max(margin1, margin2) + margin = margin1 + margin2 max_x1 = -MJ_MAXVAL max_y1 = -MJ_MAXVAL @@ -249,10 +227,7 @@ def _obb_filter( xmat2: wp.mat33, ) -> bool: """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB.""" - if BLEEDING_EDGE_MUJOCO: - margin = margin1 + margin2 - else: - margin = wp.max(margin1, margin2) + margin = margin1 + margin2 xcenter = mat23() normal = mat63() @@ -315,13 +290,25 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound # 8: obb aabb_id = worldid % ngeom_aabb if wp.static(ngeom_aabb > 1) else 0 - center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0] - size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1] + center1, center2 = ( + geom_aabb[aabb_id, geom1, 0], + geom_aabb[aabb_id, geom2, 0], + ) # kernel_analyzer: ignore + size1, size2 = ( + geom_aabb[aabb_id, geom1, 1], + geom_aabb[aabb_id, geom2, 1], + ) # kernel_analyzer: ignore rbound_id = worldid % ngeom_rbound if wp.static(ngeom_rbound > 1) else 0 - rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2] + rbound1, rbound2 = ( + geom_rbound[rbound_id, geom1], + geom_rbound[rbound_id, geom2], + ) # kernel_analyzer: ignore margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0 - margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] + margin1, margin2 = ( + geom_margin[margin_id, geom1], + geom_margin[margin_id, geom2], + ) # kernel_analyzer: ignore xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2] xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2] @@ -804,3 +791,6 @@ def collision(m: Model, d: Data): sap_broadphase(m, d, ctx) _narrowphase(m, d, ctx) + + if m.callback.contactfilter: + m.callback.contactfilter(m, d) 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 b44264cd..7e4948ed 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 @@ -16,12 +16,11 @@ import math from typing import Tuple -import warp as wp - -from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 +import warp as wp # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) @@ -581,14 +580,13 @@ def gjk( simplex_index1 = wp.vec4i() simplex_index2 = wp.vec4i() n = int(0) - cnt = int(1) coordinates = wp.vec4() # barycentric coordinates epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance) # set initial guess x_k = x1_0 - x2_0 - for _ in range(gjk_iterations): + for k in range(gjk_iterations): xnorm = wp.dot(x_k, x_k) # TODO(kbayes): determine new constant here if xnorm < 1e-12: @@ -665,10 +663,12 @@ def gjk( if n == 4: break - cnt += 1 - - if cnt == gjk_iterations: - wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations) + if k == gjk_iterations - 1: + wp.printf( + "Warning: opt.ccd_iterations, currently set to %d, needs to be" + " increased.\n", + gjk_iterations, + ) result = GJKResult() @@ -1220,14 +1220,13 @@ def _epa( idx = int(-1) pidx = int(-1) epsilon = wp.where(is_discrete, 1e-15, tolerance) - cnt = int(1) nvalid = pt.nface # number of potential faces for expanding the polytope # the face vertices are encoded in 10-bits that index the vertex array, # so iterations must be cap to limit the number of generated vertices # (one new vertex per iteration) epa_iterations = wp.min(epa_iterations, 1000) - for _ in range(epa_iterations): + for k in range(epa_iterations): pidx = idx idx = int(-1) lower2 = float(FLOAT_MAX) @@ -1325,10 +1324,13 @@ def _epa( # clear horizon pt.nhorizon = 0 - cnt += 1 - if cnt == epa_iterations: - wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations) + if k == epa_iterations - 1: + wp.printf( + "Warning: opt.ccd_iterations, currently set to %d, needs to be" + " increased.\n", + gjk_iterations, + ) # return from valid face if idx > -1: 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 67df5e1a..c7a40514 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 @@ -15,8 +15,11 @@ from typing import Tuple -import warp as wp - +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import box_box from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import capsule_box from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import capsule_capsule @@ -29,142 +32,21 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_capsule from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_cylinder from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_sphere -from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame -from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL -from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext -from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType -from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import mat43 -from mujoco.mjx.third_party.mujoco_warp._src.types import mat63 +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL +from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) -@wp.struct -class Geom: - pos: wp.vec3 - rot: wp.mat33 - normal: wp.vec3 - size: wp.vec3 - margin: float - hfprism: mat63 - vertadr: int - vertnum: int - vert: wp.array(dtype=wp.vec3) - graphadr: int - graph: wp.array(dtype=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) - 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), - # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - # In: - geoms: wp.vec2i, - worldid: int, -) -> Tuple[Geom, Geom]: - geom1 = Geom() - geom2 = Geom() - - g1 = geoms[0] - g2 = geoms[1] - geom_type1 = geom_type[g1] - geom_type2 = geom_type[g2] - - geom1.pos = geom_xpos_in[worldid, g1] - geom1.rot = geom_xmat_in[worldid, g1] - geom1.size = geom_size[worldid % geom_size.shape[0], g1] - geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2]) # plane - - geom2.pos = geom_xpos_in[worldid, g2] - geom2.rot = geom_xmat_in[worldid, g2] - geom2.size = geom_size[worldid % geom_size.shape[0], g2] - geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2]) # plane - - if geom_type1 == GeomType.MESH: - dataid = geom_dataid[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) - geom1.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1) - geom1.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1) - - geom1.vert = mesh_vert - geom1.graph = mesh_graph - geom1.mesh_polynormal = mesh_polynormal - geom1.mesh_polyvertadr = mesh_polyvertadr - geom1.mesh_polyvertnum = mesh_polyvertnum - geom1.mesh_polyvert = mesh_polyvert - geom1.mesh_polymapadr = mesh_polymapadr - geom1.mesh_polymapnum = mesh_polymapnum - geom1.mesh_polymap = mesh_polymap - - if geom_type2 == GeomType.MESH: - dataid = geom_dataid[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) - geom2.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1) - geom2.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1) - - geom2.vert = mesh_vert - geom2.graph = mesh_graph - geom2.mesh_polynormal = mesh_polynormal - geom2.mesh_polyvertadr = mesh_polyvertadr - geom2.mesh_polyvertnum = mesh_polyvertnum - geom2.mesh_polyvert = mesh_polyvert - geom2.mesh_polymapadr = mesh_polymapadr - geom2.mesh_polymapnum = mesh_polymapnum - geom2.mesh_polymap = mesh_polymap - - geom1.index = -1 - geom1.margin = 0.0 - - geom2.index = -1 - geom2.margin = 0.0 - - return geom1, geom2 - - @wp.func def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43, wp.vec3]: """Core contact geometry calculation for plane-convex collision. @@ -394,183 +276,6 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup return contact_dist, contact_pos, plane_normal -@wp.func -def write_contact( - # Data in: - naconmax_in: int, - # In: - id_: int, - dist_in: float, - pos_in: wp.vec3, - frame_in: wp.mat33, - margin_in: float, - gap_in: float, - condim_in: int, - friction_in: vec5, - solref_in: wp.vec2, - solreffriction_in: wp.vec2, - solimp_in: vec5, - geoms_in: wp.vec2i, - pairid_in: wp.vec2i, - worldid_in: int, - # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), -) -> int: - active = dist_in < margin_in - - # skip contact and no collision sensor - if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1: - return 0 - - contact_type = 0 - - if pairid_in[0] >= -1 and active: - contact_type |= ContactType.CONSTRAINT - - if pairid_in[1] >= 0: - contact_type |= ContactType.SENSOR - - cid = wp.atomic_add(nacon_out, 0, 1) - if cid < naconmax_in: - contact_dist_out[cid] = dist_in - contact_pos_out[cid] = pos_in - contact_frame_out[cid] = frame_in - contact_geom_out[cid] = geoms_in - contact_worldid_out[cid] = worldid_in - includemargin = margin_in - gap_in - contact_includemargin_out[cid] = includemargin - contact_dim_out[cid] = condim_in - contact_friction_out[cid] = friction_in - contact_solref_out[cid] = solref_in - contact_solreffriction_out[cid] = solreffriction_in - contact_solimp_out[cid] = solimp_in - contact_type_out[cid] = contact_type - contact_geomcollisionid_out[cid] = id_ - return int(active) - return 0 - - -@wp.func -def contact_params( - # Model: - geom_condim: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), - # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), - cid: int, - worldid: int, -): - geoms = collision_pair_in[cid] - pairid = collision_pairid_in[cid][0] - - # TODO(team): early return if collision sensor but no contact - # (ie, pairid[0] < -1 and pairid[1] < 0) - - if pairid > -1: - margin = pair_margin[worldid, pairid] - gap = pair_gap[worldid, pairid] - condim = pair_dim[pairid] - friction = pair_friction[worldid, pairid] - solref = pair_solref[worldid, pairid] - solreffriction = pair_solreffriction[worldid, pairid] - solimp = pair_solimp[worldid, pairid] - else: - g1 = geoms[0] - g2 = geoms[1] - solmix_id = worldid % geom_solmix.shape[0] - friction_id = worldid % geom_friction.shape[0] - solref_id = worldid % geom_solref.shape[0] - solimp_id = worldid % geom_solimp.shape[0] - margin_id = worldid % geom_margin.shape[0] - gap_id = worldid % geom_gap.shape[0] - - solmix1 = geom_solmix[solmix_id, g1] - solmix2 = geom_solmix[solmix_id, g2] - - condim1 = geom_condim[g1] - condim2 = geom_condim[g2] - - # priority - p1 = geom_priority[g1] - p2 = geom_priority[g2] - - if p1 > p2: - mix = 1.0 - condim = condim1 - max_geom_friction = geom_friction[friction_id, g1] - elif p2 > p1: - mix = 0.0 - condim = condim2 - max_geom_friction = geom_friction[friction_id, g2] - else: - mix = safe_div(solmix1, solmix1 + solmix2) - mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix) - mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix) - mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix) - condim = wp.max(condim1, condim2) - max_geom_friction = wp.max(geom_friction[friction_id, g1], geom_friction[friction_id, g2]) - - friction = vec5( - max_geom_friction[0], - max_geom_friction[0], - max_geom_friction[1], - max_geom_friction[2], - max_geom_friction[2], - ) - - if geom_solref[solref_id, g1][0] > 0.0 and geom_solref[solref_id, g2][0] > 0.0: - solref = mix * geom_solref[solref_id, g1] + (1.0 - mix) * geom_solref[solref_id, g2] - else: - solref = wp.min(geom_solref[solref_id, g1], geom_solref[solref_id, g2]) - - solreffriction = wp.vec2(0.0, 0.0) - solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2] - # geom priority is ignored - if BLEEDING_EDGE_MUJOCO: - margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2] - gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2] - else: - margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2]) - gap = wp.max(geom_gap[gap_id, g1], geom_gap[gap_id, g2]) - - friction = vec5( - wp.max(MJ_MINMU, friction[0]), - wp.max(MJ_MINMU, friction[1]), - wp.max(MJ_MINMU, friction[2]), - wp.max(MJ_MINMU, friction[3]), - wp.max(MJ_MINMU, friction[4]), - ) - - return geoms, margin, gap, condim, friction, solref, solreffriction, solimp - - @wp.func def plane_sphere_wrapper( # Data in: 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 15ea3a41..c84b67c8 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 @@ -15,14 +15,12 @@ from typing import Tuple -import warp as wp - -from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params -from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair -from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair +from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh -from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType from mujoco.mjx.third_party.mujoco_warp._src.types import Model @@ -31,6 +29,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec8 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -56,6 +55,7 @@ class VolumeData: oct_aabb: wp.array2d(dtype=wp.vec3) oct_child: wp.array(dtype=vec8i) oct_coeff: wp.array(dtype=vec8) + root: int = 0 valid: bool = False @@ -77,17 +77,18 @@ 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), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=wp.vec3f), - # In: - g_type: int, - g_size: wp.vec3, - plugin_id: int, - mesh_id: int, + # Model: + oct_child: wp.array(dtype=vec8i), + oct_aabb: wp.array2d(dtype=wp.vec3), + oct_coeff: wp.array(dtype=vec8), + mesh_octadr: wp.array(dtype=int), + plugin: wp.array(dtype=int), + plugin_attr: wp.array(dtype=wp.vec3f), + # In: + g_type: int, + g_size: wp.vec3, + plugin_id: int, + mesh_id: int, ) -> Tuple[wp.vec3, int, VolumeData, MeshData]: attributes = g_size plugin_index = -1 @@ -98,8 +99,10 @@ def get_sdf_params( plugin_index = plugin[plugin_id] elif g_type == GeomType.SDF and mesh_id != -1: - volume_data.center = oct_aabb[mesh_id, 0] - volume_data.half_size = oct_aabb[mesh_id, 1] + octadr = mesh_octadr[mesh_id] + volume_data.center = oct_aabb[octadr, 0] + volume_data.half_size = oct_aabb[octadr, 1] + volume_data.root = octadr volume_data.oct_aabb = oct_aabb volume_data.oct_child = oct_child volume_data.oct_coeff = oct_coeff @@ -225,9 +228,13 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3: @wp.func def find_oct( - oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool + oct_child: wp.array(dtype=vec8i), + oct_aabb: wp.array2d(dtype=wp.vec3), + p: wp.vec3, + grad: bool, + root: int, ) -> Tuple[int, Tuple[vec8, vec8, vec8]]: - stack = int(0) + stack = root niter = int(100) rx = vec8(0.0) ry = vec8(0.0) @@ -258,15 +265,17 @@ def find_oct( coord = wp.cw_div(p - vmin, vmax - vmin) # check if the node is a leaf + # child indices are relative to root (mesh_octadr offset) + child0 = oct_child[node][0] if ( - oct_child[node][0] == -1 - and oct_child[node][1] == -1 - and oct_child[node][2] == -1 - and oct_child[node][3] == -1 - and oct_child[node][4] == -1 - and oct_child[node][5] == -1 - and oct_child[node][6] == -1 - and oct_child[node][7] == -1 + child0 == -1 + and oct_child[node][1] == -1 + and oct_child[node][2] == -1 + and oct_child[node][3] == -1 + and oct_child[node][4] == -1 + and oct_child[node][5] == -1 + and oct_child[node][6] == -1 + and oct_child[node][7] == -1 ): for j in range(8): if not grad: @@ -282,10 +291,12 @@ def find_oct( return node, (rx, ry, rz) # compute which of 8 children to visit next + # child indices are stored relative to mesh_octadr, add root offset x = 0 if coord[0] < 0.5 else 1 y = 0 if coord[1] < 0.5 else 1 z = 0 if coord[2] < 0.5 else 1 - stack = oct_child[node][4 * z + 2 * y + x] + child = oct_child[node][4 * z + 2 * y + x] + stack = child + root if child != -1 else -1 wp.print("ERROR: Node not found\n") return -1, (rx, ry, rz) @@ -331,7 +342,13 @@ def box_project(center: wp.vec3, half_size: wp.vec3, xyz: wp.vec3) -> Tuple[floa @wp.func def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float: dist0, point = box_project(volume_data.center, volume_data.half_size, xyz) - node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=False) + node, weights = find_oct( + volume_data.oct_child, + volume_data.oct_aabb, + point, + grad=False, + root=volume_data.root, + ) return dist0 + wp.dot(weights[0], volume_data.oct_coeff[node]) @@ -348,7 +365,13 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3: grad_y = (sample_volume_sdf(xyz + dy, volume_data) - f) / h grad_z = (sample_volume_sdf(xyz + dz, volume_data) - f) / h return wp.vec3(grad_x, grad_y, grad_z) - node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=True) + node, weights = find_oct( + volume_data.oct_child, + volume_data.oct_aabb, + point, + grad=True, + root=volume_data.root, + ) grad_x = wp.dot(weights[0], volume_data.oct_coeff[node]) grad_y = wp.dot(weights[1], volume_data.oct_coeff[node]) grad_z = wp.dot(weights[2], volume_data.oct_coeff[node]) @@ -622,75 +645,76 @@ def gradient_descent( @wp.kernel def _sdf_narrowphase( - # Model: - nmeshface: int, - oct_child: wp.array(dtype=vec8i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8), - geom_type: wp.array(dtype=int), - geom_condim: wp.array(dtype=int), - geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array2d(dtype=wp.vec3), - geom_aabb: wp.array3d(dtype=wp.vec3), - geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array2d(dtype=float), - mesh_vertadr: wp.array(dtype=int), - mesh_vertnum: wp.array(dtype=int), - mesh_faceadr: wp.array(dtype=int), - mesh_graphadr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_face: wp.array(dtype=wp.vec3i), - mesh_graph: wp.array(dtype=int), - mesh_polynum: wp.array(dtype=int), - mesh_polyadr: wp.array(dtype=int), - mesh_polynormal: wp.array(dtype=wp.vec3), - mesh_polyvertadr: wp.array(dtype=int), - mesh_polyvertnum: wp.array(dtype=int), - mesh_polyvert: wp.array(dtype=int), - mesh_polymapadr: wp.array(dtype=int), - mesh_polymapnum: wp.array(dtype=int), - mesh_polymap: wp.array(dtype=int), - pair_dim: wp.array(dtype=int), - pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array2d(dtype=vec5), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=wp.vec3f), - geom_plugin_index: wp.array(dtype=int), - # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - naconmax_in: int, - ncollision_in: wp.array(dtype=int), - # In: - collision_pair_in: wp.array(dtype=wp.vec2i), - collision_pairid_in: wp.array(dtype=wp.vec2i), - collision_worldid_in: wp.array(dtype=int), - sdf_initpoints: int, - sdf_iterations: int, - # Data out: - contact_dist_out: wp.array(dtype=float), - contact_pos_out: wp.array(dtype=wp.vec3), - contact_frame_out: wp.array(dtype=wp.mat33), - contact_includemargin_out: wp.array(dtype=float), - contact_friction_out: wp.array(dtype=vec5), - contact_solref_out: wp.array(dtype=wp.vec2), - contact_solreffriction_out: wp.array(dtype=wp.vec2), - contact_solimp_out: wp.array(dtype=vec5), - contact_dim_out: wp.array(dtype=int), - contact_geom_out: wp.array(dtype=wp.vec2i), - contact_worldid_out: wp.array(dtype=int), - contact_type_out: wp.array(dtype=int), - contact_geomcollisionid_out: wp.array(dtype=int), - nacon_out: wp.array(dtype=int), + # Model: + nmeshface: int, + oct_child: wp.array(dtype=vec8i), + oct_aabb: wp.array2d(dtype=wp.vec3), + oct_coeff: wp.array(dtype=vec8), + geom_type: wp.array(dtype=int), + geom_condim: wp.array(dtype=int), + geom_dataid: wp.array(dtype=int), + geom_priority: wp.array(dtype=int), + geom_solmix: wp.array2d(dtype=float), + geom_solref: wp.array2d(dtype=wp.vec2), + geom_solimp: wp.array2d(dtype=vec5), + geom_size: wp.array2d(dtype=wp.vec3), + geom_aabb: wp.array3d(dtype=wp.vec3), + geom_friction: wp.array2d(dtype=wp.vec3), + geom_margin: wp.array2d(dtype=float), + geom_gap: wp.array2d(dtype=float), + mesh_vertadr: wp.array(dtype=int), + mesh_vertnum: wp.array(dtype=int), + mesh_faceadr: wp.array(dtype=int), + mesh_octadr: wp.array(dtype=int), + mesh_graphadr: wp.array(dtype=int), + mesh_vert: wp.array(dtype=wp.vec3), + mesh_face: wp.array(dtype=wp.vec3i), + mesh_graph: wp.array(dtype=int), + mesh_polynum: wp.array(dtype=int), + mesh_polyadr: wp.array(dtype=int), + mesh_polynormal: wp.array(dtype=wp.vec3), + mesh_polyvertadr: wp.array(dtype=int), + mesh_polyvertnum: wp.array(dtype=int), + mesh_polyvert: wp.array(dtype=int), + mesh_polymapadr: wp.array(dtype=int), + mesh_polymapnum: wp.array(dtype=int), + mesh_polymap: wp.array(dtype=int), + pair_dim: wp.array(dtype=int), + pair_solref: wp.array2d(dtype=wp.vec2), + pair_solreffriction: wp.array2d(dtype=wp.vec2), + pair_solimp: wp.array2d(dtype=vec5), + pair_margin: wp.array2d(dtype=float), + pair_gap: wp.array2d(dtype=float), + pair_friction: wp.array2d(dtype=vec5), + plugin: wp.array(dtype=int), + plugin_attr: wp.array(dtype=wp.vec3f), + geom_plugin_index: wp.array(dtype=int), + # Data in: + geom_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xmat_in: wp.array2d(dtype=wp.mat33), + naconmax_in: int, + ncollision_in: wp.array(dtype=int), + # In: + collision_pair_in: wp.array(dtype=wp.vec2i), + collision_pairid_in: wp.array(dtype=wp.vec2i), + collision_worldid_in: wp.array(dtype=int), + sdf_initpoints: int, + sdf_iterations: int, + # Data out: + contact_dist_out: wp.array(dtype=float), + contact_pos_out: wp.array(dtype=wp.vec3), + contact_frame_out: wp.array(dtype=wp.mat33), + contact_includemargin_out: wp.array(dtype=float), + contact_friction_out: wp.array(dtype=vec5), + contact_solref_out: wp.array(dtype=wp.vec2), + contact_solreffriction_out: wp.array(dtype=wp.vec2), + contact_solimp_out: wp.array(dtype=vec5), + contact_dim_out: wp.array(dtype=int), + contact_geom_out: wp.array(dtype=wp.vec2i), + contact_worldid_out: wp.array(dtype=int), + contact_type_out: wp.array(dtype=int), + contact_geomcollisionid_out: wp.array(dtype=int), + nacon_out: wp.array(dtype=int), ): i, contact_tid = wp.tid() if i >= sdf_initpoints: @@ -775,11 +799,29 @@ def _sdf_narrowphase( rot1 = geom1.rot attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params( - oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1] + oct_child, + oct_aabb, + oct_coeff, + mesh_octadr, + plugin, + plugin_attr, + type1, + geom1.size, + g1_plugin, + geom_dataid[g1], ) attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params( - oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2] + oct_child, + oct_aabb, + oct_coeff, + mesh_octadr, + plugin, + plugin_attr, + type2, + geom2.size, + g2_plugin, + geom_dataid[g2], ) mesh_data1.nmeshface = nmeshface @@ -868,75 +910,76 @@ def _sdf_narrowphase( @event_scope def sdf_narrowphase(m: Model, d: Data, ctx: CollisionContext): wp.launch( - _sdf_narrowphase, - dim=(m.opt.sdf_initpoints, d.naconmax), - inputs=[ - m.nmeshface, - m.oct_child, - m.oct_aabb, - m.oct_coeff, - m.geom_type, - m.geom_condim, - m.geom_dataid, - m.geom_priority, - m.geom_solmix, - m.geom_solref, - m.geom_solimp, - m.geom_size, - m.geom_aabb, - m.geom_friction, - m.geom_margin, - m.geom_gap, - m.mesh_vertadr, - m.mesh_vertnum, - m.mesh_faceadr, - m.mesh_graphadr, - m.mesh_vert, - m.mesh_face, - m.mesh_graph, - m.mesh_polynum, - m.mesh_polyadr, - m.mesh_polynormal, - m.mesh_polyvertadr, - m.mesh_polyvertnum, - m.mesh_polyvert, - m.mesh_polymapadr, - m.mesh_polymapnum, - m.mesh_polymap, - m.pair_dim, - m.pair_solref, - m.pair_solreffriction, - m.pair_solimp, - m.pair_margin, - m.pair_gap, - m.pair_friction, - m.plugin, - m.plugin_attr, - m.geom_plugin_index, - d.geom_xpos, - d.geom_xmat, - d.naconmax, - d.ncollision, - ctx.collision_pair, - ctx.collision_pairid, - ctx.collision_worldid, - m.opt.sdf_initpoints, - m.opt.sdf_iterations, - ], - outputs=[ - d.contact.dist, - d.contact.pos, - d.contact.frame, - d.contact.includemargin, - d.contact.friction, - d.contact.solref, - d.contact.solreffriction, - d.contact.solimp, - d.contact.dim, - d.contact.geom, - d.contact.worldid, - d.contact.type, - d.contact.geomcollisionid, - d.nacon, - ], + _sdf_narrowphase, + dim=(m.opt.sdf_initpoints, d.naconmax), + inputs=[ + m.nmeshface, + m.oct_child, + m.oct_aabb, + m.oct_coeff, + m.geom_type, + m.geom_condim, + m.geom_dataid, + m.geom_priority, + m.geom_solmix, + m.geom_solref, + m.geom_solimp, + m.geom_size, + m.geom_aabb, + m.geom_friction, + m.geom_margin, + m.geom_gap, + m.mesh_vertadr, + m.mesh_vertnum, + m.mesh_faceadr, + m.mesh_octadr, + m.mesh_graphadr, + m.mesh_vert, + m.mesh_face, + m.mesh_graph, + m.mesh_polynum, + m.mesh_polyadr, + m.mesh_polynormal, + m.mesh_polyvertadr, + m.mesh_polyvertnum, + m.mesh_polyvert, + m.mesh_polymapadr, + m.mesh_polymapnum, + m.mesh_polymap, + m.pair_dim, + m.pair_solref, + m.pair_solreffriction, + m.pair_solimp, + m.pair_margin, + m.pair_gap, + m.pair_friction, + m.plugin, + m.plugin_attr, + m.geom_plugin_index, + d.geom_xpos, + d.geom_xmat, + d.naconmax, + d.ncollision, + ctx.collision_pair, + ctx.collision_pairid, + ctx.collision_worldid, + m.opt.sdf_initpoints, + m.opt.sdf_iterations, + ], + outputs=[ + d.contact.dist, + d.contact.pos, + d.contact.frame, + d.contact.includemargin, + d.contact.friction, + d.contact.solref, + d.contact.solreffriction, + d.contact.solimp, + d.contact.dim, + d.contact.geom, + d.contact.worldid, + d.contact.type, + d.contact.geomcollisionid, + d.nacon, + ], ) 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 f786e36d..228006e8 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py @@ -13,16 +13,18 @@ # limitations under the License. # ============================================================================== -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import support from mujoco.mjx.third_party.mujoco_warp._src import types from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType -from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 +from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit +from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN from mujoco.mjx.third_party.mujoco_warp._src.types import vec11 +from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -45,33 +47,34 @@ def _zero_constraint_counts( @wp.func -def _update_efc_row( - # In: - worldid: int, - timestep: float, - refsafe: int, - efcid: int, - pos_aref: float, - pos_imp: float, - invweight: float, - solref: wp.vec2, - solimp: vec5, - margin: float, - vel: float, - frictionloss: float, - type: int, - id: int, - # Data out: - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - 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), +def _efc_row( + # Model: + opt_disableflags: int, + # In: + worldid: int, + timestep: float, + efcid: int, + pos_aref: float, + pos_imp: float, + invweight: float, + solref: wp.vec2, + solimp: vec5, + margin: float, + vel: float, + frictionloss: float, + 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), ): - # Calculate kbi + # calculate kbi timeconst = solref[0] dampratio = solref[1] dmin = solimp[0] @@ -80,8 +83,7 @@ def _update_efc_row( mid = solimp[3] power = solimp[4] - # TODO(team): wp.static? - if not refsafe: + if not (opt_disableflags & DisableBit.REFSAFE): timeconst = wp.max(timeconst, 2.0 * timestep) dmin = wp.clamp(dmin, types.MJ_MINIMP, types.MJ_MAXIMP) @@ -90,10 +92,11 @@ def _update_efc_row( mid = wp.clamp(mid, types.MJ_MINIMP, types.MJ_MAXIMP) power = wp.max(1.0, power) - # See https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters - k = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) + # see https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters + dmax_sq = dmax * dmax + k = 1.0 / (dmax_sq * timeconst * timeconst * dampratio * dampratio) b = 2.0 / (dmax * timeconst) - k = wp.where(solref[0] <= 0, -solref[0] / (dmax * dmax), k) + k = wp.where(solref[0] <= 0, -solref[0] / dmax_sq, k) b = wp.where(solref[1] <= 0, -solref[1] / dmax, b) imp_x = wp.abs(pos_imp) / width @@ -104,58 +107,67 @@ def _update_efc_row( imp = wp.clamp(imp, dmin, dmax) imp = wp.where(imp_x > 1.0, dmax, imp) - # Update constraints - efc_D_out[worldid, efcid] = 1.0 / wp.max(invweight * (1.0 - imp) / imp, types.MJ_MINVAL) - efc_vel_out[worldid, efcid] = vel - efc_aref_out[worldid, efcid] = -k * imp * pos_aref - b * vel - efc_pos_out[worldid, efcid] = pos_aref + margin - efc_margin_out[worldid, efcid] = margin - efc_frictionloss_out[worldid, efcid] = frictionloss - efc_type_out[worldid, efcid] = type - efc_id_out[worldid, efcid] = id + # set outputs + D_out[worldid, efcid] = 1.0 / wp.max( + invweight * (1.0 - imp) / imp, types.MJ_MINVAL + ) + vel_out[worldid, efcid] = vel + aref_out[worldid, efcid] = -k * imp * pos_aref - b * vel + pos_out[worldid, efcid] = pos_aref + margin + margin_out[worldid, efcid] = margin + frictionloss_out[worldid, efcid] = frictionloss + type_out[worldid, efcid] = type + id_out[worldid, efcid] = id @wp.kernel -def _efc_equality_connect( - # Model: - nv: int, - nsite: int, - opt_timestep: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: 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), - eq_connect_adr: wp.array(dtype=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), - njmax_in: int, - # In: - refsafe_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_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), +def _equality_connect( + # Model: + nv: int, + nsite: int, + opt_timestep: wp.array(dtype=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), + is_sparse: bool, + eq_connect_adr: wp.array(dtype=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), + njmax_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), ): """Calculates constraint rows for connect equality constraints.""" worldid, eqconnectid = wp.tid() @@ -167,7 +179,7 @@ def _efc_equality_connect( wp.atomic_add(ne_out, worldid, 3) efcid = wp.atomic_add(nefc_out, worldid, 3) - if efcid + 3 >= njmax_in: + if efcid >= njmax_in - 3: return data = eq_data[worldid % eq_data.shape[0], eqid] @@ -177,54 +189,132 @@ def _efc_equality_connect( obj1id = eq_obj1id[eqid] obj2id = eq_obj2id[eqid] - if nsite and eq_objtype[eqid] == types.ObjType.SITE: - # body1id stores the index of site_bodyid. - body1id = site_bodyid[obj1id] - body2id = site_bodyid[obj2id] + if nsite > 0 and eq_objtype[eqid] == types.ObjType.SITE: + body1 = site_bodyid[obj1id] + body2 = site_bodyid[obj2id] pos1 = site_xpos_in[worldid, obj1id] pos2 = site_xpos_in[worldid, obj2id] else: - body1id = obj1id - body2id = obj2id - pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor1 - pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor2 + body1 = obj1id + body2 = obj2id + pos1 = xpos_in[worldid, body1] + xmat_in[worldid, body1] @ anchor1 + pos2 = xpos_in[worldid, body2] + xmat_in[worldid, body2] @ anchor2 # error is difference in global positions pos = pos1 - pos2 # compute Jacobian difference (opposite of contact: 0 - 1) Jqvel = wp.vec3f(0.0, 0.0, 0.0) - for dofid in range(nv): # TODO: parallelize - jacp1, _ = support.jac_dof( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos1, - body1id, - dofid, - worldid, - ) - jacp2, _ = support.jac_dof( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos2, - body2id, - dofid, - worldid, - ) - j1mj2 = jacp1 - jacp2 - efc_J_out[worldid, efcid + 0, dofid] = j1mj2[0] - efc_J_out[worldid, efcid + 1, dofid] = j1mj2[1] - efc_J_out[worldid, efcid + 2, dofid] = j1mj2[2] - Jqvel += j1mj2 * qvel_in[worldid, dofid] + + if is_sparse: + body1 = body_weldid[body1] + body2 = body_weldid[body2] + + da1 = int(body_dofadr[body1] + body_dofnum[body1] - 1) + da2 = int(body_dofadr[body2] + body_dofnum[body2] - 1) + + efcid0 = efcid + 0 + efcid1 = efcid + 1 + efcid2 = efcid + 2 + + rowadr0 = efcid0 * nv + rowadr1 = efcid1 * nv + rowadr2 = efcid2 * nv + + efc_J_rowadr_out[worldid, efcid0] = rowadr0 + efc_J_rowadr_out[worldid, efcid1] = rowadr1 + efc_J_rowadr_out[worldid, efcid2] = rowadr2 + + rownnz = int(0) + + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + if da1 == da: + da1 = dof_parentid[da1] + if da2 == da: + da2 = dof_parentid[da2] + + jacp1, _ = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + da, + worldid, + ) + jacp2, _ = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + da, + worldid, + ) + j1mj2 = jacp1 - jacp2 + + sparseid0 = rowadr0 + rownnz + sparseid1 = rowadr1 + rownnz + sparseid2 = rowadr2 + rownnz + + efc_J_colind_out[worldid, 0, sparseid0] = da + efc_J_colind_out[worldid, 0, sparseid1] = da + efc_J_colind_out[worldid, 0, sparseid2] = da + + efc_J_out[worldid, 0, sparseid0] = j1mj2[0] + efc_J_out[worldid, 0, sparseid1] = j1mj2[1] + efc_J_out[worldid, 0, sparseid2] = j1mj2[2] + + Jqvel += j1mj2 * qvel_in[worldid, da] + + rownnz += 1 + + efc_J_rownnz_out[worldid, efcid0] = rownnz + efc_J_rownnz_out[worldid, efcid1] = rownnz + efc_J_rownnz_out[worldid, efcid2] = rownnz + else: + # TODO(team): dof tree traversal + for dofid in range(nv): + jacp1, _ = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + dofid, + worldid, + ) + jacp2, _ = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + dofid, + worldid, + ) + j1mj2 = jacp1 - jacp2 + + efc_J_out[worldid, efcid + 0, dofid] = j1mj2[0] + efc_J_out[worldid, efcid + 1, dofid] = j1mj2[1] + efc_J_out[worldid, efcid + 2, dofid] = j1mj2[2] + + Jqvel += j1mj2 * qvel_in[worldid, dofid] body_invweight0_id = worldid % body_invweight0.shape[0] - invweight = body_invweight0[body_invweight0_id, body1id][0] + body_invweight0[body_invweight0_id, body2id][0] + invweight = ( + body_invweight0[body_invweight0_id, body1][0] + + body_invweight0[body_invweight0_id, body2][0] + ) pos_imp = wp.length(pos) solref = eq_solref[worldid % eq_solref.shape[0], eqid] @@ -234,66 +324,69 @@ def _efc_equality_connect( for i in range(3): efcidi = efcid + i - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcidi, - pos[i], - pos_imp, - invweight, - solref, - solimp, - 0.0, - Jqvel[i], - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + timestep, + efcidi, + pos[i], + pos_imp, + invweight, + solref, + solimp, + 0.0, + Jqvel[i], + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_equality_joint( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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), - eq_jnt_adr: wp.array(dtype=int), - # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - njmax_in: int, - # In: - refsafe_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_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), +def _equality_joint( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + is_sparse: bool, + eq_jnt_adr: wp.array(dtype=int), + # Data in: + qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), + eq_active_in: wp.array2d(dtype=bool), + njmax_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), ): worldid, eqjntid = wp.tid() eqid = eq_jnt_adr[eqjntid] @@ -307,18 +400,29 @@ def _efc_equality_joint( if efcid >= njmax_in: return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - jntid_1 = eq_obj1id[eqid] jntid_2 = eq_obj2id[eqid] data = eq_data[worldid % eq_data.shape[0], eqid] dofadr1 = jnt_dofadr[jntid_1] qposadr1 = jnt_qposadr[jntid_1] - efc_J_out[worldid, efcid, dofadr1] = 1.0 qpos0_id = worldid % qpos0.shape[0] dof_invweight0_id = worldid % dof_invweight0.shape[0] + if is_sparse: + if jntid_2 > -1: + rownnz = 2 + else: + rownnz = 1 + efc_J_rownnz_out[worldid, efcid] = rownnz + rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = rowadr + efc_J_colind_out[worldid, 0, rowadr] = dofadr1 + efc_J_out[worldid, 0, rowadr] = 1.0 + else: + for i in range(nv): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dofadr1] = 1.0 + if jntid_2 > -1: # Two joint constraint qposadr2 = jnt_qposadr[jntid_2] @@ -326,14 +430,26 @@ def _efc_equality_joint( dif = qpos_in[worldid, qposadr2] - qpos0[qpos0_id, qposadr2] # Horner's method for polynomials - rhs = data[0] + dif * (data[1] + dif * (data[2] + dif * (data[3] + dif * data[4]))) - deriv_2 = data[1] + dif * (2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4])) + rhs = data[0] + dif * ( + data[1] + dif * (data[2] + dif * (data[3] + dif * data[4])) + ) + deriv_2 = data[1] + dif * ( + 2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4]) + ) pos = qpos_in[worldid, qposadr1] - qpos0[qpos0_id, qposadr1] - rhs Jqvel = qvel_in[worldid, dofadr1] - qvel_in[worldid, dofadr2] * deriv_2 - invweight = dof_invweight0[dof_invweight0_id, dofadr1] + dof_invweight0[dof_invweight0_id, dofadr2] + invweight = ( + dof_invweight0[dof_invweight0_id, dofadr1] + + dof_invweight0[dof_invweight0_id, dofadr2] + ) - efc_J_out[worldid, efcid, dofadr2] = -deriv_2 + if is_sparse: + sparseid = rowadr + 1 + efc_J_colind_out[worldid, 0, sparseid] = dofadr2 + efc_J_out[worldid, 0, sparseid] = -deriv_2 + else: + efc_J_out[worldid, efcid, dofadr2] = -deriv_2 else: # Single joint constraint pos = qpos_in[worldid, qposadr1] - qpos0[qpos0_id, qposadr1] - data[0] @@ -341,65 +457,68 @@ def _efc_equality_joint( invweight = dof_invweight0[dof_invweight0_id, dofadr1] # Update constraint parameters - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - invweight, - eq_solref[worldid % eq_solref.shape[0], eqid], - eq_solimp[worldid % eq_solimp.shape[0], eqid], - 0.0, - Jqvel, - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + invweight, + eq_solref[worldid % eq_solref.shape[0], eqid], + eq_solimp[worldid % eq_solimp.shape[0], eqid], + 0.0, + Jqvel, + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_equality_tendon( - # Model: - nv: int, - opt_timestep: wp.array(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), - tendon_length0: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), - eq_ten_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - ten_J_in: wp.array3d(dtype=float), - ten_length_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _equality_tendon( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + tendon_length0: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), + is_sparse: bool, + eq_ten_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + eq_active_in: wp.array2d(dtype=bool), + ten_J_in: wp.array3d(dtype=float), + ten_length_in: wp.array2d(dtype=float), + njmax_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), ): worldid, eqtenid = wp.tid() eqid = eq_ten_adr[eqtenid] @@ -421,13 +540,21 @@ def _efc_equality_tendon( solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] tendon_length0_id = worldid % tendon_length0.shape[0] tendon_invweight0_id = worldid % tendon_invweight0.shape[0] - pos1 = ten_length_in[worldid, obj1id] - tendon_length0[tendon_length0_id, obj1id] + pos1 = ( + ten_length_in[worldid, obj1id] - tendon_length0[tendon_length0_id, obj1id] + ) jac1 = ten_J_in[worldid, obj1id] if obj2id > -1: - invweight = tendon_invweight0[tendon_invweight0_id, obj1id] + tendon_invweight0[tendon_invweight0_id, obj2id] + invweight = ( + tendon_invweight0[tendon_invweight0_id, obj1id] + + tendon_invweight0[tendon_invweight0_id, obj2id] + ) - pos2 = ten_length_in[worldid, obj2id] - tendon_length0[tendon_length0_id, obj2id] + pos2 = ( + ten_length_in[worldid, obj2id] + - tendon_length0[tendon_length0_id, obj2id] + ) jac2 = ten_J_in[worldid, obj2id] dif = pos2 @@ -435,162 +562,520 @@ def _efc_equality_tendon( dif3 = dif2 * dif dif4 = dif3 * dif - pos = pos1 - (data[0] + data[1] * dif + data[2] * dif2 + data[3] * dif3 + data[4] * dif4) - deriv = data[1] + 2.0 * data[2] * dif + 3.0 * data[3] * dif2 + 4.0 * data[4] * dif3 + pos = pos1 - ( + data[0] + + data[1] * dif + + data[2] * dif2 + + data[3] * dif3 + + data[4] * dif4 + ) + deriv = ( + data[1] + + 2.0 * data[2] * dif + + 3.0 * data[3] * dif2 + + 4.0 * data[4] * dif3 + ) else: invweight = tendon_invweight0[tendon_invweight0_id, obj1id] pos = pos1 - data[0] deriv = 0.0 Jqvel = float(0.0) + + # TODO(team): sparse tendon jacobian + if is_sparse: + rowadr = efcid * nv + efc_J_rownnz_out[worldid, efcid] = nv + efc_J_rowadr_out[worldid, efcid] = rowadr + for i in range(nv): if deriv != 0.0: J = jac1[i] + jac2[i] * -deriv else: J = jac1[i] - efc_J_out[worldid, efcid, i] = J + if is_sparse: + efc_J_colind_out[worldid, 0, rowadr + i] = i + efc_J_out[worldid, 0, rowadr + i] = J + else: + efc_J_out[worldid, efcid, i] = J Jqvel += J * qvel_in[worldid, i] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - invweight, - solref, - solimp, - 0.0, - Jqvel, - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + invweight, + solref, + solimp, + 0.0, + Jqvel, + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) +@cache_kernel +def _equality_flex(is_sparse: bool): + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # Model: + nv: int, + opt_timestep: wp.array(dtype=float), + opt_disableflags: 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_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_flex_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + flexedge_J_in: wp.array2d(dtype=float), + flexedge_length_in: wp.array2d(dtype=float), + njmax_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), + ): + worldid, eqflexid, edgeid = wp.tid() + eqid = eq_flex_adr[eqflexid] + + wp.atomic_add(ne_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= njmax_in: + return + + pos = flexedge_length_in[worldid, edgeid] - flexedge_length0[edgeid] + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + + Jqvel = float(0.0) + + rownnz = flexedge_J_rownnz[edgeid] + flex_rowadr = flexedge_J_rowadr[edgeid] + + if wp.static(is_sparse): + efc_J_rownnz_out[worldid, efcid] = rownnz + efc_rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = efc_rowadr + for i in range(rownnz): + flex_sparseid = flex_rowadr + i + efc_sparseid = efc_rowadr + i + colind = flexedge_J_colind[flex_sparseid] + J = flexedge_J_in[worldid, flex_sparseid] + efc_J_colind_out[worldid, 0, efc_sparseid] = colind + efc_J_out[worldid, 0, efc_sparseid] = J + Jqvel += J * qvel_in[worldid, colind] + else: + for i in range(nv): + efc_J_out[worldid, efcid, i] = 0.0 + for i in range(rownnz): + flex_sparseid = flex_rowadr + i + colind = flexedge_J_colind[flex_sparseid] + J = flexedge_J_in[worldid, flex_sparseid] + efc_J_out[worldid, efcid, colind] = J + Jqvel += J * qvel_in[worldid, colind] + + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + flexedge_invweight0[edgeid], + solref, + solimp, + 0.0, + Jqvel, + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) + + return kernel + + @wp.kernel -def _efc_equality_flex( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_flex_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - flexedge_J_in: wp.array3d(dtype=float), - flexedge_length_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _equality_weld( + # Model: + nv: int, + nsite: int, + opt_timestep: wp.array(dtype=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), + is_sparse: bool, + eq_wld_adr: wp.array(dtype=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), + njmax_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), ): - worldid, eqflexid, edgeid = wp.tid() - eqid = eq_flex_adr[eqflexid] + worldid, eqweldid = wp.tid() + eqid = eq_wld_adr[eqweldid] - wp.atomic_add(ne_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) - - if efcid >= njmax_in: + if not eq_active_in[worldid, eqid]: return - pos = flexedge_length_in[worldid, edgeid] - flexedge_length0[edgeid] + wp.atomic_add(ne_out, worldid, 6) + efcid = wp.atomic_add(nefc_out, worldid, 6) + + if efcid >= njmax_in - 6: + return + + is_site = eq_objtype[eqid] == types.ObjType.SITE and nsite > 0 + + obj1id = eq_obj1id[eqid] + obj2id = eq_obj2id[eqid] + + data = eq_data[worldid % eq_data.shape[0], eqid] + anchor1 = wp.vec3(data[0], data[1], data[2]) + anchor2 = wp.vec3(data[3], data[4], data[5]) + relpose = wp.quat(data[6], data[7], data[8], data[9]) + torquescale = data[10] + + if is_site: + body1 = site_bodyid[obj1id] + body2 = site_bodyid[obj2id] + pos1 = site_xpos_in[worldid, obj1id] + pos2 = site_xpos_in[worldid, obj2id] + + site_quat_id = worldid % site_quat.shape[0] + quat = math.mul_quat( + xquat_in[worldid, body1], site_quat[site_quat_id, obj1id] + ) + quat1 = math.quat_inv( + math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id]) + ) + + else: + body1 = obj1id + body2 = obj2id + pos1 = xpos_in[worldid, body1] + xmat_in[worldid, body1] @ anchor2 + pos2 = xpos_in[worldid, body2] + xmat_in[worldid, body2] @ anchor1 + + quat = math.mul_quat(xquat_in[worldid, body1], relpose) + quat1 = math.quat_inv(xquat_in[worldid, body2]) + + # compute Jacobian difference (opposite of contact: 0 - 1) + Jqvelp = wp.vec3f(0.0, 0.0, 0.0) + Jqvelr = wp.vec3f(0.0, 0.0, 0.0) + + if is_sparse: + body1 = body_weldid[body1] + body2 = body_weldid[body2] + + da1 = int(body_dofadr[body1] + body_dofnum[body1] - 1) + da2 = int(body_dofadr[body2] + body_dofnum[body2] - 1) + + efcid0 = efcid + 0 + efcid1 = efcid + 1 + efcid2 = efcid + 2 + efcid3 = efcid + 3 + efcid4 = efcid + 4 + efcid5 = efcid + 5 + + rowadr0 = efcid0 * nv + rowadr1 = efcid1 * nv + rowadr2 = efcid2 * nv + rowadr3 = efcid3 * nv + rowadr4 = efcid4 * nv + rowadr5 = efcid5 * nv + + efc_J_rowadr_out[worldid, efcid0] = rowadr0 + efc_J_rowadr_out[worldid, efcid1] = rowadr1 + efc_J_rowadr_out[worldid, efcid2] = rowadr2 + efc_J_rowadr_out[worldid, efcid3] = rowadr3 + efc_J_rowadr_out[worldid, efcid4] = rowadr4 + efc_J_rowadr_out[worldid, efcid5] = rowadr5 + + rownnz = int(0) + + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + if da1 == da: + da1 = dof_parentid[da] + if da2 == da: + da2 = dof_parentid[da] + + jacp1, jacr1 = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + da, + worldid, + ) + jacp2, jacr2 = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + da, + worldid, + ) + + jacdifp = jacp1 - jacp2 + + jacdifr = (jacr1 - jacr2) * torquescale + jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat) + jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3]) + + sparseid0 = rowadr0 + rownnz + sparseid1 = rowadr1 + rownnz + sparseid2 = rowadr2 + rownnz + sparseid3 = rowadr3 + rownnz + sparseid4 = rowadr4 + rownnz + sparseid5 = rowadr5 + rownnz + + efc_J_colind_out[worldid, 0, sparseid0] = da + efc_J_colind_out[worldid, 0, sparseid1] = da + efc_J_colind_out[worldid, 0, sparseid2] = da + efc_J_colind_out[worldid, 0, sparseid3] = da + efc_J_colind_out[worldid, 0, sparseid4] = da + efc_J_colind_out[worldid, 0, sparseid5] = da + + efc_J_out[worldid, 0, sparseid0] = jacdifp[0] + efc_J_out[worldid, 0, sparseid1] = jacdifp[1] + efc_J_out[worldid, 0, sparseid2] = jacdifp[2] + efc_J_out[worldid, 0, sparseid3] = jacdifr[0] + efc_J_out[worldid, 0, sparseid4] = jacdifr[1] + efc_J_out[worldid, 0, sparseid5] = jacdifr[2] + + Jqvelp += jacdifp * qvel_in[worldid, da] + Jqvelr += jacdifr * qvel_in[worldid, da] + + rownnz += 1 + + efc_J_rownnz_out[worldid, efcid0] = rownnz + efc_J_rownnz_out[worldid, efcid1] = rownnz + efc_J_rownnz_out[worldid, efcid2] = rownnz + efc_J_rownnz_out[worldid, efcid3] = rownnz + efc_J_rownnz_out[worldid, efcid4] = rownnz + efc_J_rownnz_out[worldid, efcid5] = rownnz + else: + for dofid in range(nv): + jacp1, jacr1 = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + dofid, + worldid, + ) + jacp2, jacr2 = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + dofid, + worldid, + ) + + jacdifp = jacp1 - jacp2 + + for i in range(3): + efc_J_out[worldid, efcid + i, dofid] = jacdifp[i] + + jacdifr = (jacr1 - jacr2) * torquescale + jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat) + jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3]) + + for i in range(3): + efc_J_out[worldid, efcid + 3 + i, dofid] = jacdifr[i] + + Jqvelp += jacdifp * qvel_in[worldid, dofid] + Jqvelr += jacdifr * qvel_in[worldid, dofid] + + # error is difference in global position and orientation + cpos = pos1 - pos2 + + crotq = math.mul_quat(quat1, quat) # copy axis components + crot = wp.vec3(crotq[1], crotq[2], crotq[3]) * torquescale + + body_invweight0_id = worldid % body_invweight0.shape[0] + invweight_t = ( + body_invweight0[body_invweight0_id, body1][0] + + body_invweight0[body_invweight0_id, body2][0] + ) + + pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) + solref = eq_solref[worldid % eq_solref.shape[0], eqid] solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] - Jqvel = float(0.0) + timestep = opt_timestep[worldid % opt_timestep.shape[0]] - # TODO(team): remove once efc.J is sparse - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 + for i in range(3): + _efc_row( + opt_disableflags, + worldid, + timestep, + efcid + i, + cpos[i], + pos_imp, + invweight_t, + solref, + solimp, + 0.0, + Jqvelp[i], + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) - rownnz = flexedge_J_rownnz[edgeid] - rowadr = flexedge_J_rowadr[edgeid] - for i in range(rownnz): - sparseid = rowadr + i - colind = flexedge_J_colind[sparseid] - J = flexedge_J_in[worldid, 0, sparseid] - # TODO(team): sparse efc.J - efc_J_out[worldid, efcid, colind] = J - Jqvel += J * qvel_in[worldid, colind] - - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - flexedge_invweight0[edgeid], - solref, - solimp, - 0.0, - Jqvel, - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + invweight_r = ( + body_invweight0[body_invweight0_id, body1][1] + + body_invweight0[body_invweight0_id, body2][1] ) + for i in range(3): + _efc_row( + opt_disableflags, + worldid, + timestep, + efcid + 3 + i, + crot[i], + pos_imp, + invweight_r, + solref, + solimp, + 0.0, + Jqvelr[i], + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) + @wp.kernel -def _efc_friction_dof( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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), - # Data in: - qvel_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _friction_dof( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + is_sparse: bool, + # Data in: + qvel_in: wp.array2d(dtype=float), + njmax_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), ): worldid, dofid = wp.tid() @@ -605,68 +1090,78 @@ def _efc_friction_dof( if efcid >= njmax_in: return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 + if is_sparse: + efc_J_rownnz_out[worldid, efcid] = 1 + rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = rowadr + efc_J_colind_out[worldid, 0, rowadr] = dofid + efc_J_out[worldid, 0, rowadr] = 1.0 + else: + for i in range(nv): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dofid] = 1.0 - efc_J_out[worldid, efcid, dofid] = 1.0 Jqvel = qvel_in[worldid, dofid] dof_invweight0_id = worldid % dof_invweight0.shape[0] dof_solref_id = worldid % dof_solref.shape[0] dof_solimp_id = worldid % dof_solimp.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - 0.0, - 0.0, - dof_invweight0[dof_invweight0_id, dofid], - dof_solref[dof_solref_id, dofid], - dof_solimp[dof_solimp_id, dofid], - 0.0, - Jqvel, - dof_frictionloss[dof_frictionloss_id, dofid], - ConstraintType.FRICTION_DOF, - dofid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + 0.0, + 0.0, + dof_invweight0[dof_invweight0_id, dofid], + dof_solref[dof_solref_id, dofid], + dof_solimp[dof_solimp_id, dofid], + 0.0, + Jqvel, + dof_frictionloss[dof_frictionloss_id, dofid], + ConstraintType.FRICTION_DOF, + dofid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_friction_tendon( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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), - # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array3d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _friction_tendon( + # Model: + nv: int, + opt_timestep: wp.array(dtype=float), + opt_disableflags: 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), + is_sparse: bool, + # Data in: + qvel_in: wp.array2d(dtype=float), + ten_J_in: wp.array3d(dtype=float), + njmax_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), ): worldid, tenid = wp.tid() @@ -684,272 +1179,86 @@ def _efc_friction_tendon( Jqvel = float(0.0) - # TODO(team): parallelize + # TODO(team): sparse tendon jacobian + if is_sparse: + rowadr = efcid * nv + efc_J_rownnz_out[worldid, efcid] = nv + efc_J_rowadr_out[worldid, efcid] = rowadr + for i in range(nv): + # TODO(team): sparse ten_J J = ten_J_in[worldid, tenid, i] - efc_J_out[worldid, efcid, i] = J + if is_sparse: + efc_J_colind_out[worldid, 0, rowadr + i] = i + efc_J_out[worldid, 0, rowadr + i] = J + else: + efc_J_out[worldid, efcid, i] = J + Jqvel += J * qvel_in[worldid, i] tendon_invweight0_id = worldid % tendon_invweight0.shape[0] tendon_solref_fri_id = worldid % tendon_solref_fri.shape[0] tendon_solimp_fri_id = worldid % tendon_solimp_fri.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - 0.0, - 0.0, - tendon_invweight0[tendon_invweight0_id, tenid], - tendon_solref_fri[tendon_solref_fri_id, tenid], - tendon_solimp_fri[tendon_solimp_fri_id, tenid], - 0.0, - Jqvel, - frictionloss, - ConstraintType.FRICTION_TENDON, - tenid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + 0.0, + 0.0, + tendon_invweight0[tendon_invweight0_id, tenid], + tendon_solref_fri[tendon_solref_fri_id, tenid], + tendon_solimp_fri[tendon_solimp_fri_id, tenid], + 0.0, + Jqvel, + frictionloss, + ConstraintType.FRICTION_TENDON, + tenid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_equality_weld( - # Model: - nv: int, - nsite: int, - opt_timestep: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: 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), - eq_wld_adr: wp.array(dtype=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), - njmax_in: int, - # In: - refsafe_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_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), -): - worldid, eqweldid = wp.tid() - eqid = eq_wld_adr[eqweldid] - - if not eq_active_in[worldid, eqid]: - return - - wp.atomic_add(ne_out, worldid, 6) - efcid = wp.atomic_add(nefc_out, worldid, 6) - - if efcid + 6 >= njmax_in: - return - - is_site = eq_objtype[eqid] == types.ObjType.SITE and nsite > 0 - - obj1id = eq_obj1id[eqid] - obj2id = eq_obj2id[eqid] - - data = eq_data[worldid % eq_data.shape[0], eqid] - anchor1 = wp.vec3(data[0], data[1], data[2]) - anchor2 = wp.vec3(data[3], data[4], data[5]) - relpose = wp.quat(data[6], data[7], data[8], data[9]) - torquescale = data[10] - - if is_site: - # body1id stores the index of site_bodyid. - body1id = site_bodyid[obj1id] - body2id = site_bodyid[obj2id] - pos1 = site_xpos_in[worldid, obj1id] - pos2 = site_xpos_in[worldid, obj2id] - - site_quat_id = worldid % site_quat.shape[0] - quat = math.mul_quat(xquat_in[worldid, body1id], site_quat[site_quat_id, obj1id]) - quat1 = math.quat_inv(math.mul_quat(xquat_in[worldid, body2id], site_quat[site_quat_id, obj2id])) - - else: - body1id = obj1id - body2id = obj2id - pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor2 - pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor1 - - quat = math.mul_quat(xquat_in[worldid, body1id], relpose) - quat1 = math.quat_inv(xquat_in[worldid, body2id]) - - # compute Jacobian difference (opposite of contact: 0 - 1) - Jqvelp = wp.vec3f(0.0, 0.0, 0.0) - Jqvelr = wp.vec3f(0.0, 0.0, 0.0) - - for dofid in range(nv): # TODO: parallelize - jacp1, jacr1 = support.jac_dof( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos1, - body1id, - dofid, - worldid, - ) - jacp2, jacr2 = support.jac_dof( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos2, - body2id, - dofid, - worldid, - ) - - jacdifp = jacp1 - jacp2 - for i in range(3): - efc_J_out[worldid, efcid + i, dofid] = jacdifp[i] - - jacdifr = (jacr1 - jacr2) * torquescale - jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat) - jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3]) - - for i in range(3): - efc_J_out[worldid, efcid + 3 + i, dofid] = jacdifr[i] - - Jqvelp += jacdifp * qvel_in[worldid, dofid] - Jqvelr += jacdifr * qvel_in[worldid, dofid] - - # error is difference in global position and orientation - cpos = pos1 - pos2 - - crotq = math.mul_quat(quat1, quat) # copy axis components - crot = wp.vec3(crotq[1], crotq[2], crotq[3]) * torquescale - - body_invweight0_id = worldid % body_invweight0.shape[0] - invweight_t = body_invweight0[body_invweight0_id, body1id][0] + body_invweight0[body_invweight0_id, body2id][0] - - pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) - - solref = eq_solref[worldid % eq_solref.shape[0], eqid] - solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] - - timestep = opt_timestep[worldid % opt_timestep.shape[0]] - - for i in range(3): - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcid + i, - cpos[i], - pos_imp, - invweight_t, - solref, - solimp, - 0.0, - Jqvelp[i], - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) - - invweight_r = body_invweight0[body_invweight0_id, body1id][1] + body_invweight0[body_invweight0_id, body2id][1] - - for i in range(3): - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcid + 3 + i, - crot[i], - pos_imp, - invweight_r, - solref, - solimp, - 0.0, - Jqvelr[i], - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) - - -@wp.kernel -def _efc_limit_slide_hinge( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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_limited_slide_hinge_adr: wp.array(dtype=int), - # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _limit_slide_hinge( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + is_sparse: bool, + jnt_limited_slide_hinge_adr: wp.array(dtype=int), + # Data in: + qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), + njmax_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), ): worldid, jntlimitedid = wp.tid() jntid = jnt_limited_slide_hinge_adr[jntlimitedid] @@ -970,75 +1279,86 @@ def _efc_limit_slide_hinge( if efcid >= njmax_in: return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - dofadr = jnt_dofadr[jntid] J = float(dist_min < dist_max) * 2.0 - 1.0 - efc_J_out[worldid, efcid, dofadr] = J + + if is_sparse: + efc_J_rownnz_out[worldid, efcid] = 1 + rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = rowadr + efc_J_colind_out[worldid, 0, rowadr] = dofadr + efc_J_out[worldid, 0, rowadr] = J + else: + for i in range(nv): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dofadr] = J + Jqvel = J * qvel_in[worldid, dofadr] dof_invweight0_id = worldid % dof_invweight0.shape[0] jnt_solref_id = worldid % jnt_solref.shape[0] jnt_solimp_id = worldid % jnt_solimp.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - dof_invweight0[dof_invweight0_id, dofadr], - jnt_solref[jnt_solref_id, jntid], - jnt_solimp[jnt_solimp_id, jntid], - jntmargin, - Jqvel, - 0.0, - ConstraintType.LIMIT_JOINT, - jntid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + dof_invweight0[dof_invweight0_id, dofadr], + jnt_solref[jnt_solref_id, jntid], + jnt_solimp[jnt_solimp_id, jntid], + jntmargin, + Jqvel, + 0.0, + ConstraintType.LIMIT_JOINT, + jntid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_limit_ball( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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_limited_ball_adr: wp.array(dtype=int), - # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _limit_ball( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + is_sparse: bool, + jnt_limited_ball_adr: wp.array(dtype=int), + # Data in: + qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), + njmax_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), ): worldid, jntlimitedid = wp.tid() jntid = jnt_limited_ball_adr[jntlimitedid] @@ -1064,83 +1384,105 @@ def _efc_limit_ball( if efcid >= njmax_in: return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - dofadr = jnt_dofadr[jntid] + dof0 = dofadr + 0 + dof1 = dofadr + 1 + dof2 = dofadr + 2 - efc_J_out[worldid, efcid, dofadr + 0] = -axis[0] - efc_J_out[worldid, efcid, dofadr + 1] = -axis[1] - efc_J_out[worldid, efcid, dofadr + 2] = -axis[2] + if is_sparse: + efc_J_rownnz_out[worldid, efcid] = 3 + rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = rowadr - Jqvel = -axis[0] * qvel_in[worldid, dofadr + 0] - Jqvel -= axis[1] * qvel_in[worldid, dofadr + 1] - Jqvel -= axis[2] * qvel_in[worldid, dofadr + 2] + sparseid0 = rowadr + 0 + sparseid1 = rowadr + 1 + sparseid2 = rowadr + 2 + + efc_J_colind_out[worldid, 0, sparseid0] = dof0 + efc_J_colind_out[worldid, 0, sparseid1] = dof1 + efc_J_colind_out[worldid, 0, sparseid2] = dof2 + + efc_J_out[worldid, 0, sparseid0] = -axis[0] + efc_J_out[worldid, 0, sparseid1] = -axis[1] + efc_J_out[worldid, 0, sparseid2] = -axis[2] + else: + for i in range(nv): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dof0] = -axis[0] + efc_J_out[worldid, efcid, dof1] = -axis[1] + efc_J_out[worldid, efcid, dof2] = -axis[2] + + Jqvel = -axis[0] * qvel_in[worldid, dof0] + Jqvel -= axis[1] * qvel_in[worldid, dof1] + Jqvel -= axis[2] * qvel_in[worldid, dof2] dof_invweight0_id = worldid % dof_invweight0.shape[0] jnt_solref_id = worldid % jnt_solref.shape[0] jnt_solimp_id = worldid % jnt_solimp.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - dof_invweight0[dof_invweight0_id, dofadr], - jnt_solref[jnt_solref_id, jntid], - jnt_solimp[jnt_solimp_id, jntid], - jntmargin, - Jqvel, - 0.0, - ConstraintType.LIMIT_JOINT, - jntid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + dof_invweight0[dof_invweight0_id, dofadr], + jnt_solref[jnt_solref_id, jntid], + jnt_solimp[jnt_solimp_id, jntid], + jntmargin, + Jqvel, + 0.0, + ConstraintType.LIMIT_JOINT, + jntid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_limit_tendon( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - jnt_dofadr: wp.array(dtype=int), - tendon_adr: wp.array(dtype=int), - tendon_num: 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), - wrap_type: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - tendon_limited_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array3d(dtype=float), - ten_length_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_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_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), +def _limit_tendon( + # Model: + nv: int, + opt_timestep: wp.array(dtype=float), + opt_disableflags: int, + jnt_dofadr: wp.array(dtype=int), + tendon_adr: wp.array(dtype=int), + tendon_num: 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), + wrap_type: wp.array(dtype=int), + wrap_objid: wp.array(dtype=int), + is_sparse: bool, + tendon_limited_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + ten_J_in: wp.array3d(dtype=float), + ten_length_in: wp.array2d(dtype=float), + njmax_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), ): worldid, tenlimitedid = wp.tid() tenid = tendon_limited_adr[tenlimitedid] @@ -1161,101 +1503,125 @@ def _efc_limit_tendon( if efcid >= njmax_in: return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - Jqvel = float(0.0) scl = float(dist_min < dist_max) * 2.0 - 1.0 + # TODO(team): sparse tendon jacobian + if is_sparse: + rowadr = efcid * nv + efc_J_rownnz_out[worldid, efcid] = nv + efc_J_rowadr_out[worldid, efcid] = rowadr + for i in range(nv): + efc_J_colind_out[worldid, 0, rowadr + i] = i + efc_J_out[worldid, 0, rowadr + i] = 0.0 + adr = tendon_adr[tenid] if wrap_type[adr] == types.WrapType.JOINT: + if not is_sparse: + for i in range(nv): + efc_J_out[worldid, efcid, i] = 0.0 + ten_num = tendon_num[tenid] for i in range(ten_num): dofadr = jnt_dofadr[wrap_objid[adr + i]] J = scl * ten_J_in[worldid, tenid, dofadr] - efc_J_out[worldid, efcid, dofadr] = J + + if is_sparse: + efc_J_out[worldid, 0, rowadr + dofadr] = J + else: + efc_J_out[worldid, efcid, dofadr] = J + Jqvel += J * qvel_in[worldid, dofadr] else: for i in range(nv): J = scl * ten_J_in[worldid, tenid, i] - efc_J_out[worldid, efcid, i] = J + + if is_sparse: + efc_J_out[worldid, 0, rowadr + i] = J + else: + efc_J_out[worldid, efcid, i] = J + Jqvel += J * qvel_in[worldid, i] tendon_invweight0_id = worldid % tendon_invweight0.shape[0] tendon_solref_lim_id = worldid % tendon_solref_lim.shape[0] tendon_solimp_lim_id = worldid % tendon_solimp_lim.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - tendon_invweight0[tendon_invweight0_id, tenid], - tendon_solref_lim[tendon_solref_lim_id, tenid], - tendon_solimp_lim[tendon_solimp_lim_id, tenid], - tenmargin, - Jqvel, - 0.0, - ConstraintType.LIMIT_TENDON, - tenid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + tendon_invweight0[tendon_invweight0_id, tenid], + tendon_solref_lim[tendon_solref_lim_id, tenid], + tendon_solimp_lim[tendon_solimp_lim_id, tenid], + tenmargin, + Jqvel, + 0.0, + ConstraintType.LIMIT_TENDON, + tenid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_contact_pyramidal( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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), - # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - njmax_in: int, - nacon_in: wp.array(dtype=int), - # In: - refsafe_in: int, - 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), - 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), - # 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_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), +def _contact_pyramidal( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + 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), + njmax_in: int, + nacon_in: wp.array(dtype=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), + 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), + # 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), ): conid, dimid = wp.tid() @@ -1308,6 +1674,10 @@ def _efc_contact_pyramidal( invweight = invweight + fri0 * fri0 * invweight invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt + if is_sparse: + rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = rowadr + Jqvel = float(0.0) # skip fixed bodies @@ -1318,7 +1688,20 @@ def _efc_contact_pyramidal( da2 = body_dofadr[body2] + body_dofnum[body2] - 1 da = wp.max(da1, da2) - for dofid in range(nv - 1, -1, -1): + if is_sparse: + rownnz = int(0) + dofid = int(da) + else: + dofid = int(nv - 1) + + while True: + if is_sparse: + if da1 < 0 and da2 < 0: + break + else: + if dofid < 0: + break + if dofid == da: # TODO(team): contact_jacobian jac1p, jac1r = support.jac_dof( @@ -1365,7 +1748,13 @@ def _efc_contact_pyramidal( else: J -= Ji * frii - efc_J_out[worldid, efcid, dofid] = J + if is_sparse: + sparseid = rowadr + rownnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + rownnz += 1 + else: + efc_J_out[worldid, efcid, dofid] = J Jqvel += J * qvel_in[worldid, dofid] # Advance tree pointers and recompute da for next iteration @@ -1374,87 +1763,100 @@ def _efc_contact_pyramidal( if da2 == da: da2 = dof_parentid[da2] da = wp.max(da1, da2) + if is_sparse: + dofid = da + else: + dofid -= 1 else: - efc_J_out[worldid, efcid, dofid] = 0.0 + if not is_sparse: + efc_J_out[worldid, efcid, dofid] = 0.0 + dofid -= 1 + + if is_sparse: + efc_J_rownnz_out[worldid, efcid] = rownnz if condim == 1: efc_type = ConstraintType.CONTACT_FRICTIONLESS else: efc_type = ConstraintType.CONTACT_PYRAMIDAL - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcid, - pos, - pos, - invweight, - solref_in[conid], - solimp_in[conid], - includemargin, - Jqvel, - 0.0, - efc_type, - conid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + timestep, + efcid, + pos, + pos, + invweight, + solref_in[conid], + solimp_in[conid], + includemargin, + Jqvel, + 0.0, + efc_type, + conid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @wp.kernel -def _efc_contact_elliptic( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - 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), - # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - njmax_in: int, - nacon_in: wp.array(dtype=int), - # In: - refsafe_in: int, - 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), - 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), - # 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_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), +def _contact_elliptic( + # Model: + nv: int, + opt_timestep: wp.array(dtype=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), + 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), + njmax_in: int, + nacon_in: wp.array(dtype=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), + 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), + # 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), ): conid, dimid = wp.tid() @@ -1485,13 +1887,17 @@ def _efc_contact_elliptic( impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] contact_efc_address_out[conid, dimid] = efcid - con_pos = pos_in[conid] - frame = frame_in[conid] - geom = geom_in[conid] body1 = geom_bodyid[geom[0]] body2 = geom_bodyid[geom[1]] + con_pos = pos_in[conid] + frame = frame_in[conid] + + if is_sparse: + rowadr = efcid * nv + efc_J_rowadr_out[worldid, efcid] = rowadr + Jqvel = float(0.0) # skip fixed bodies @@ -1502,7 +1908,20 @@ def _efc_contact_elliptic( da2 = body_dofadr[body2] + body_dofnum[body2] - 1 da = wp.max(da1, da2) - for dofid in range(nv - 1, -1, -1): + if is_sparse: + rownnz = int(0) + dofid = int(da) + else: + dofid = int(nv - 1) + + while True: + if is_sparse: + if da1 < 0 and da2 < 0: + break + else: + if dofid < 0: + break + if dofid == da: # TODO(team): contact jacobian jac1p, jac1r = support.jac_dof( @@ -1537,7 +1956,13 @@ def _efc_contact_elliptic( jac_dif = jac2r[xyz] - jac1r[xyz] J += frame[dimid - 3, xyz] * jac_dif - efc_J_out[worldid, efcid, dofid] = J + if is_sparse: + sparseid = rowadr + rownnz + efc_J_colind_out[worldid, 0, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = J + rownnz += 1 + else: + efc_J_out[worldid, efcid, dofid] = J Jqvel += J * qvel_in[worldid, dofid] # Advance tree pointers and recompute da for next iteration @@ -1546,8 +1971,17 @@ def _efc_contact_elliptic( if da2 == da: da2 = dof_parentid[da2] da = wp.max(da1, da2) + if is_sparse: + dofid = da + else: + dofid -= 1 else: - efc_J_out[worldid, efcid, dofid] = 0.0 + if not is_sparse: + efc_J_out[worldid, efcid, dofid] = 0.0 + dofid -= 1 + + if is_sparse: + efc_J_rownnz_out[worldid, efcid] = rownnz body_invweight0_id = worldid % body_invweight0.shape[0] invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] @@ -1578,29 +2012,29 @@ def _efc_contact_elliptic( else: efc_type = ConstraintType.CONTACT_ELLIPTIC - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcid, - pos_aref, - pos, - invweight, - ref, - solimp_in[conid], - includemargin, - Jqvel, - 0.0, - efc_type, - conid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, + _efc_row( + opt_disableflags, + worldid, + timestep, + efcid, + pos_aref, + pos, + invweight, + ref, + solimp_in[conid], + includemargin, + Jqvel, + 0.0, + efc_type, + conid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, ) @@ -1613,471 +2047,527 @@ def make_constraint(m: types.Model, d: types.Data): inputs=[d.ne, d.nf, d.nl, d.nefc], ) - if not (m.opt.disableflags & types.DisableBit.CONSTRAINT): - refsafe = m.opt.disableflags & types.DisableBit.REFSAFE + if types.SPARSE_CONSTRAINT_JACOBIAN: + d.contact.efc_address.fill_(-1) + if not (m.opt.disableflags & types.DisableBit.CONSTRAINT): if not (m.opt.disableflags & types.DisableBit.EQUALITY): wp.launch( - _efc_equality_connect, - dim=(d.nworld, m.eq_connect_adr.size), - inputs=[ - m.nv, - m.nsite, - m.opt.timestep, - m.body_parentid, - m.body_rootid, - m.body_invweight0, - m.dof_bodyid, - m.site_bodyid, - m.eq_obj1id, - m.eq_obj2id, - m.eq_objtype, - m.eq_solref, - m.eq_solimp, - m.eq_data, - m.eq_connect_adr, - d.qvel, - d.eq_active, - d.xpos, - d.xmat, - d.site_xpos, - d.subtree_com, - d.cdof, - d.njmax, - refsafe, - ], - outputs=[ - d.ne, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _equality_connect, + dim=(d.nworld, m.eq_connect_adr.size), + inputs=[ + m.nv, + m.nsite, + m.opt.timestep, + m.opt.disableflags, + m.body_parentid, + m.body_rootid, + m.body_weldid, + m.body_dofnum, + m.body_dofadr, + m.body_invweight0, + m.dof_bodyid, + m.dof_parentid, + m.site_bodyid, + m.eq_obj1id, + m.eq_obj2id, + m.eq_objtype, + m.eq_solref, + m.eq_solimp, + m.eq_data, + SPARSE_CONSTRAINT_JACOBIAN, + m.eq_connect_adr, + d.qvel, + d.eq_active, + d.xpos, + d.xmat, + d.site_xpos, + d.subtree_com, + d.cdof, + d.njmax, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_equality_weld, - dim=(d.nworld, m.eq_wld_adr.size), - inputs=[ - m.nv, - m.nsite, - m.opt.timestep, - m.body_parentid, - m.body_rootid, - m.body_invweight0, - m.dof_bodyid, - m.site_bodyid, - m.site_quat, - m.eq_obj1id, - m.eq_obj2id, - m.eq_objtype, - m.eq_solref, - m.eq_solimp, - m.eq_data, - m.eq_wld_adr, - d.qvel, - d.eq_active, - d.xpos, - d.xquat, - d.xmat, - d.site_xpos, - d.subtree_com, - d.cdof, - d.njmax, - refsafe, - ], - outputs=[ - d.ne, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _equality_weld, + dim=(d.nworld, m.eq_wld_adr.size), + inputs=[ + m.nv, + m.nsite, + m.opt.timestep, + m.opt.disableflags, + m.body_parentid, + m.body_rootid, + m.body_weldid, + m.body_dofnum, + m.body_dofadr, + m.body_invweight0, + m.dof_bodyid, + m.dof_parentid, + m.site_bodyid, + m.site_quat, + m.eq_obj1id, + m.eq_obj2id, + m.eq_objtype, + m.eq_solref, + m.eq_solimp, + m.eq_data, + SPARSE_CONSTRAINT_JACOBIAN, + m.eq_wld_adr, + d.qvel, + d.eq_active, + d.xpos, + d.xquat, + d.xmat, + d.site_xpos, + d.subtree_com, + d.cdof, + d.njmax, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_equality_joint, - dim=(d.nworld, m.eq_jnt_adr.size), - inputs=[ - m.nv, - m.opt.timestep, - m.qpos0, - m.jnt_qposadr, - m.jnt_dofadr, - m.dof_invweight0, - m.eq_obj1id, - m.eq_obj2id, - m.eq_solref, - m.eq_solimp, - m.eq_data, - m.eq_jnt_adr, - d.qpos, - d.qvel, - d.eq_active, - d.njmax, - refsafe, - ], - outputs=[ - d.ne, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _equality_joint, + dim=(d.nworld, m.eq_jnt_adr.size), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.qpos0, + m.jnt_qposadr, + m.jnt_dofadr, + m.dof_invweight0, + m.eq_obj1id, + m.eq_obj2id, + m.eq_solref, + m.eq_solimp, + m.eq_data, + SPARSE_CONSTRAINT_JACOBIAN, + m.eq_jnt_adr, + d.qpos, + d.qvel, + d.eq_active, + d.njmax, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_equality_tendon, - dim=(d.nworld, m.eq_ten_adr.size), - inputs=[ - m.nv, - m.opt.timestep, - m.eq_obj1id, - m.eq_obj2id, - m.eq_solref, - m.eq_solimp, - m.eq_data, - m.tendon_length0, - m.tendon_invweight0, - m.eq_ten_adr, - d.qvel, - d.eq_active, - d.ten_J, - d.ten_length, - d.njmax, - refsafe, - ], - outputs=[ - d.ne, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _equality_tendon, + dim=(d.nworld, m.eq_ten_adr.size), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.eq_obj1id, + m.eq_obj2id, + m.eq_solref, + m.eq_solimp, + m.eq_data, + m.tendon_length0, + m.tendon_invweight0, + SPARSE_CONSTRAINT_JACOBIAN, + m.eq_ten_adr, + d.qvel, + d.eq_active, + d.ten_J, + d.ten_length, + d.njmax, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_equality_flex, - dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge), - inputs=[ - m.nv, - m.opt.timestep, - m.flexedge_length0, - m.flexedge_invweight0, - m.flexedge_J_rownnz, - m.flexedge_J_rowadr, - m.flexedge_J_colind, - m.eq_solref, - m.eq_solimp, - m.eq_flex_adr, - d.qvel, - d.flexedge_J, - d.flexedge_length, - d.njmax, - refsafe, - ], - outputs=[ - d.ne, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _equality_flex(SPARSE_CONSTRAINT_JACOBIAN), + dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.flexedge_length0, + m.flexedge_invweight0, + m.flexedge_J_rownnz, + m.flexedge_J_rowadr, + m.flexedge_J_colind, + m.eq_solref, + m.eq_solimp, + m.eq_flex_adr, + d.qvel, + d.flexedge_J, + d.flexedge_length, + d.njmax, + ], + outputs=[ + d.ne, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS): wp.launch( - _efc_friction_dof, - dim=(d.nworld, m.nv), - inputs=[ - m.nv, - m.opt.timestep, - m.dof_solref, - m.dof_solimp, - m.dof_frictionloss, - m.dof_invweight0, - d.qvel, - d.njmax, - refsafe, - ], - outputs=[ - d.nf, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _friction_dof, + dim=(d.nworld, m.nv), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.dof_solref, + m.dof_solimp, + m.dof_frictionloss, + m.dof_invweight0, + SPARSE_CONSTRAINT_JACOBIAN, + d.qvel, + d.njmax, + ], + outputs=[ + d.nf, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_friction_tendon, - dim=(d.nworld, m.ntendon), - inputs=[ - m.nv, - m.opt.timestep, - m.tendon_solref_fri, - m.tendon_solimp_fri, - m.tendon_frictionloss, - m.tendon_invweight0, - d.qvel, - d.ten_J, - d.njmax, - refsafe, - ], - outputs=[ - d.nf, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _friction_tendon, + dim=(d.nworld, m.ntendon), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.tendon_solref_fri, + m.tendon_solimp_fri, + m.tendon_frictionloss, + m.tendon_invweight0, + SPARSE_CONSTRAINT_JACOBIAN, + d.qvel, + d.ten_J, + d.njmax, + ], + outputs=[ + d.nf, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) # limit if not (m.opt.disableflags & types.DisableBit.LIMIT): wp.launch( - _efc_limit_ball, - dim=(d.nworld, m.jnt_limited_ball_adr.size), - inputs=[ - m.nv, - m.opt.timestep, - m.jnt_qposadr, - m.jnt_dofadr, - m.jnt_solref, - m.jnt_solimp, - m.jnt_range, - m.jnt_margin, - m.dof_invweight0, - m.jnt_limited_ball_adr, - d.qpos, - d.qvel, - d.njmax, - refsafe, - ], - outputs=[ - d.nl, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _limit_ball, + dim=(d.nworld, m.jnt_limited_ball_adr.size), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.jnt_qposadr, + m.jnt_dofadr, + m.jnt_solref, + m.jnt_solimp, + m.jnt_range, + m.jnt_margin, + m.dof_invweight0, + SPARSE_CONSTRAINT_JACOBIAN, + m.jnt_limited_ball_adr, + d.qpos, + d.qvel, + d.njmax, + ], + outputs=[ + d.nl, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_limit_slide_hinge, - dim=(d.nworld, m.jnt_limited_slide_hinge_adr.size), - inputs=[ - m.nv, - m.opt.timestep, - m.jnt_qposadr, - m.jnt_dofadr, - m.jnt_solref, - m.jnt_solimp, - m.jnt_range, - m.jnt_margin, - m.dof_invweight0, - m.jnt_limited_slide_hinge_adr, - d.qpos, - d.qvel, - d.njmax, - refsafe, - ], - outputs=[ - d.nl, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _limit_slide_hinge, + dim=(d.nworld, m.jnt_limited_slide_hinge_adr.size), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.jnt_qposadr, + m.jnt_dofadr, + m.jnt_solref, + m.jnt_solimp, + m.jnt_range, + m.jnt_margin, + m.dof_invweight0, + SPARSE_CONSTRAINT_JACOBIAN, + m.jnt_limited_slide_hinge_adr, + d.qpos, + d.qvel, + d.njmax, + ], + outputs=[ + d.nl, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) wp.launch( - _efc_limit_tendon, - dim=(d.nworld, m.tendon_limited_adr.size), - inputs=[ - m.nv, - m.opt.timestep, - m.jnt_dofadr, - m.tendon_adr, - m.tendon_num, - m.tendon_solref_lim, - m.tendon_solimp_lim, - m.tendon_range, - m.tendon_margin, - m.tendon_invweight0, - m.wrap_type, - m.wrap_objid, - m.tendon_limited_adr, - d.qvel, - d.ten_J, - d.ten_length, - d.njmax, - refsafe, - ], - outputs=[ - d.nl, - d.nefc, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _limit_tendon, + dim=(d.nworld, m.tendon_limited_adr.size), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.jnt_dofadr, + m.tendon_adr, + m.tendon_num, + m.tendon_solref_lim, + m.tendon_solimp_lim, + m.tendon_range, + m.tendon_margin, + m.tendon_invweight0, + m.wrap_type, + m.wrap_objid, + SPARSE_CONSTRAINT_JACOBIAN, + m.tendon_limited_adr, + d.qvel, + d.ten_J, + d.ten_length, + d.njmax, + ], + outputs=[ + d.nl, + d.nefc, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) # contact if not (m.opt.disableflags & types.DisableBit.CONTACT): if m.opt.cone == types.ConeType.PYRAMIDAL: wp.launch( - _efc_contact_pyramidal, - dim=(d.naconmax, m.nmaxpyramid), - inputs=[ - m.nv, - m.opt.timestep, - m.opt.impratio_invsqrt, - m.body_parentid, - m.body_rootid, - m.body_weldid, - m.body_dofnum, - m.body_dofadr, - m.body_invweight0, - m.dof_bodyid, - m.dof_parentid, - m.geom_bodyid, - d.qvel, - d.subtree_com, - d.cdof, - d.njmax, - d.nacon, - refsafe, - d.contact.dist, - d.contact.dim, - d.contact.includemargin, - d.contact.worldid, - d.contact.geom, - d.contact.pos, - d.contact.frame, - d.contact.friction, - d.contact.solref, - d.contact.solimp, - d.contact.type, - ], - outputs=[ - d.nefc, - d.contact.efc_address, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _contact_pyramidal, + dim=(d.naconmax, m.nmaxpyramid), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.opt.impratio_invsqrt, + m.body_parentid, + m.body_rootid, + m.body_weldid, + m.body_dofnum, + m.body_dofadr, + m.body_invweight0, + m.dof_bodyid, + m.dof_parentid, + m.geom_bodyid, + SPARSE_CONSTRAINT_JACOBIAN, + d.qvel, + d.subtree_com, + d.cdof, + d.njmax, + d.nacon, + d.contact.dist, + d.contact.dim, + d.contact.includemargin, + d.contact.worldid, + d.contact.geom, + d.contact.pos, + d.contact.frame, + d.contact.friction, + d.contact.solref, + d.contact.solimp, + d.contact.type, + ], + outputs=[ + d.nefc, + d.contact.efc_address, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) elif m.opt.cone == types.ConeType.ELLIPTIC: wp.launch( - _efc_contact_elliptic, - dim=(d.naconmax, m.nmaxcondim), - inputs=[ - m.nv, - m.opt.timestep, - m.opt.impratio_invsqrt, - m.body_parentid, - m.body_rootid, - m.body_weldid, - m.body_dofnum, - m.body_dofadr, - m.body_invweight0, - m.dof_bodyid, - m.dof_parentid, - m.geom_bodyid, - d.qvel, - d.subtree_com, - d.cdof, - d.njmax, - d.nacon, - refsafe, - d.contact.dist, - d.contact.dim, - d.contact.includemargin, - d.contact.worldid, - d.contact.geom, - d.contact.pos, - d.contact.frame, - d.contact.friction, - d.contact.solref, - d.contact.solreffriction, - d.contact.solimp, - d.contact.type, - ], - outputs=[ - d.nefc, - d.contact.efc_address, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.pos, - d.efc.margin, - d.efc.D, - d.efc.vel, - d.efc.aref, - d.efc.frictionloss, - ], + _contact_elliptic, + dim=(d.naconmax, m.nmaxcondim), + inputs=[ + m.nv, + m.opt.timestep, + m.opt.disableflags, + m.opt.impratio_invsqrt, + m.body_parentid, + m.body_rootid, + m.body_weldid, + m.body_dofnum, + m.body_dofadr, + m.body_invweight0, + m.dof_bodyid, + m.dof_parentid, + m.geom_bodyid, + SPARSE_CONSTRAINT_JACOBIAN, + d.qvel, + d.subtree_com, + d.cdof, + d.njmax, + d.nacon, + d.contact.dist, + d.contact.dim, + d.contact.includemargin, + d.contact.worldid, + d.contact.geom, + d.contact.pos, + d.contact.frame, + d.contact.friction, + d.contact.solref, + d.contact.solreffriction, + d.contact.solimp, + d.contact.type, + ], + outputs=[ + d.nefc, + d.contact.efc_address, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.vel, + d.efc.aref, + d.efc.frictionloss, + ], ) 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 67e7aad6..406c36c4 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py @@ -21,9 +21,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import DynType from mujoco.mjx.third_party.mujoco_warp._src.types import GainType from mujoco.mjx.third_party.mujoco_warp._src.types import Model -from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f -from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope wp.set_module_options({"enable_backward": False}) @@ -96,55 +94,22 @@ def _nonzero_mask(x: float) -> float: return 0.0 -@cache_kernel -def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int): - @wp.kernel(module="unique", enable_backward=False) - def kernel( - # Data in: - actuator_moment_in: wp.array3d(dtype=float), - qM_in: wp.array3d(dtype=float), - # In: - vel_in: wp.array3d(dtype=float), - adr: wp.array(dtype=int), - # Out: - qDeriv_out: wp.array3d(dtype=float), - ): - worldid, nodeid = wp.tid() - TILE_SIZE = wp.static(tile.size) - NU = wp.static(nu) - - dofid = adr[nodeid] - vel_tile = wp.tile_load(vel_in[worldid], shape=(NU, 1), bounds_check=False) - moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, TILE_SIZE), offset=(0, dofid), bounds_check=False) - moment_weighted = wp.tile_map(wp.mul, wp.tile_broadcast(vel_tile, shape=(NU, TILE_SIZE)), moment_tile) - qderiv_tile = wp.tile_matmul(wp.tile_transpose(moment_tile), moment_weighted) - - # Mask out cross-terms for DOF pairs that are structurally zero in M - # (e.g., sibling DOFs coupled only through tendons). Without this, - # stale actuation values at sibling positions make A = M - dt*qDeriv - # non-positive-definite, causing the tiled Cholesky to produce NaN. - # Dropping these terms matches MuJoCo CPU's implicitfast approximation. - qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False) - mask_tile = wp.tile_map(_nonzero_mask, qM_tile) - qderiv_tile = wp.tile_map(wp.mul, qderiv_tile, mask_tile) - - wp.tile_store(qDeriv_out[worldid], qderiv_tile, offset=(dofid, dofid), bounds_check=False) - - return kernel - - @wp.kernel def _qderiv_actuator_passive_actuation_sparse( - # Model: - nu: int, - # Data in: - actuator_moment_in: wp.array3d(dtype=float), - # In: - vel_in: wp.array2d(dtype=float), - qMi: wp.array(dtype=int), - qMj: wp.array(dtype=int), - # Out: - qDeriv_out: wp.array3d(dtype=float), + # Model: + nu: int, + is_sparse: bool, + # Data in: + moment_rownnz_in: wp.array2d(dtype=int), + moment_rowadr_in: wp.array2d(dtype=int), + moment_colind_in: wp.array2d(dtype=int), + actuator_moment_in: wp.array2d(dtype=float), + # In: + vel_in: wp.array2d(dtype=float), + qMi: wp.array(dtype=int), + qMj: wp.array(dtype=int), + # Out: + qDeriv_out: wp.array3d(dtype=float), ): worldid, elemid = wp.tid() @@ -156,12 +121,33 @@ def _qderiv_actuator_passive_actuation_sparse( if vel == 0.0: continue - moment_i = actuator_moment_in[worldid, actid, dofiid] - moment_j = actuator_moment_in[worldid, actid, dofjid] + # TODO(team): restructure sparse version for better parallelism? + moment_i = float(0.0) + moment_j = float(0.0) + + rownnz = moment_rownnz_in[worldid, actid] + rowadr = moment_rowadr_in[worldid, actid] + for i in range(rownnz): + sparseid = rowadr + i + colind = moment_colind_in[worldid, sparseid] + if colind == dofiid: + moment_i = actuator_moment_in[worldid, sparseid] + if colind == dofjid: + moment_j = actuator_moment_in[worldid, sparseid] + if moment_i != 0.0 and moment_j != 0.0: + break + + if moment_i == 0 and moment_j == 0: + continue qderiv_contrib += moment_i * moment_j * vel - qDeriv_out[worldid, 0, elemid] = qderiv_contrib + if is_sparse: + qDeriv_out[worldid, 0, elemid] = qderiv_contrib + else: + qDeriv_out[worldid, dofiid, dofjid] = qderiv_contrib + if dofiid != dofjid: + qDeriv_out[worldid, dofjid, dofiid] = qderiv_contrib @wp.kernel @@ -277,23 +263,22 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)): ], outputs=[vel], ) - if m.is_sparse: - wp.launch( + wp.launch( _qderiv_actuator_passive_actuation_sparse, dim=(d.nworld, qMi.size), - inputs=[m.nu, d.actuator_moment, vel, qMi, qMj], + inputs=[ + m.nu, + m.is_sparse, + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + d.actuator_moment, + vel, + qMi, + qMj, + ], outputs=[out], - ) - else: - vel_3d = vel.reshape(vel.shape + (1,)) - for tile in m.qM_tiles: - wp.launch_tiled( - _qderiv_actuator_passive_actuation_dense(tile, m.nu), - dim=(d.nworld, tile.adr.size), - inputs=[d.actuator_moment, d.qM, vel_3d, tile.adr], - outputs=[out], - block_dim=m.block_dim.qderiv_actuator_dense, - ) + ) wp.launch( _qderiv_actuator_passive, dim=(d.nworld, qMi.size), 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 e345c652..0dc3de14 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -15,11 +15,10 @@ from typing import Optional -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import collision_driver from mujoco.mjx.third_party.mujoco_warp._src import constraint from mujoco.mjx.third_party.mujoco_warp._src import derivative +from mujoco.mjx.third_party.mujoco_warp._src import island from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import passive from mujoco.mjx.third_party.mujoco_warp._src import sensor @@ -27,7 +26,6 @@ from mujoco.mjx.third_party.mujoco_warp._src import smooth from mujoco.mjx.third_party.mujoco_warp._src import solver from mujoco.mjx.third_party.mujoco_warp._src import util_misc from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit @@ -36,12 +34,14 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit from mujoco.mjx.third_party.mujoco_warp._src.types import GainType from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType from mujoco.mjx.third_party.mujoco_warp._src.types import JointType +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -146,6 +146,8 @@ def _next_act( if actuator_dyntype == DynType.FILTEREXACT: tau = wp.max(MJ_MINVAL, actuator_dynprm[0]) act = act_in + act_dot_scale * act_dot_in * tau * (1.0 - wp.exp(-opt_timestep / tau)) + elif actuator_dyntype == DynType.USER: + return act_in else: act = act_in + act_dot_scale * act_dot_in * opt_timestep @@ -158,36 +160,41 @@ def _next_act( @wp.kernel def _next_activation( - # Model: - opt_timestep: wp.array(dtype=float), - actuator_dyntype: wp.array(dtype=int), - actuator_actlimited: wp.array(dtype=bool), - actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_actrange: wp.array2d(dtype=wp.vec2), - # Data in: - act_in: wp.array2d(dtype=float), - act_dot_in: wp.array2d(dtype=float), - # In: - act_dot_scale: float, - limit: bool, - # Data out: - act_out: wp.array2d(dtype=float), + # Model: + opt_timestep: wp.array(dtype=float), + actuator_dyntype: wp.array(dtype=int), + actuator_actadr: wp.array(dtype=int), + actuator_actnum: wp.array(dtype=int), + actuator_actlimited: wp.array(dtype=bool), + actuator_dynprm: wp.array2d(dtype=vec10f), + actuator_actrange: wp.array2d(dtype=wp.vec2), + # Data in: + act_in: wp.array2d(dtype=float), + act_dot_in: wp.array2d(dtype=float), + # In: + act_dot_scale: float, + limit: bool, + # Data out: + act_out: wp.array2d(dtype=float), ): - worldid, actid = wp.tid() + worldid, uid = wp.tid() opt_timestep_id = worldid % opt_timestep.shape[0] actuator_dynprm_id = worldid % actuator_dynprm.shape[0] actuator_actrange_id = worldid % actuator_actrange.shape[0] - act = _next_act( - opt_timestep[opt_timestep_id], - actuator_dyntype[actid], - actuator_dynprm[actuator_dynprm_id, actid], - actuator_actrange[actuator_actrange_id, actid], - act_in[worldid, actid], - act_dot_in[worldid, actid], - act_dot_scale, - limit and actuator_actlimited[actid], - ) - act_out[worldid, actid] = act + actadr = actuator_actadr[uid] + actnum = actuator_actnum[uid] + for j in range(actadr, actadr + actnum): + act = _next_act( + opt_timestep[opt_timestep_id], + actuator_dyntype[uid], + actuator_dynprm[actuator_dynprm_id, uid], + actuator_actrange[actuator_actrange_id, uid], + act_in[worldid, j], + act_dot_in[worldid, j], + act_dot_scale, + limit and actuator_actlimited[uid], + ) + act_out[worldid, j] = act @wp.kernel @@ -229,20 +236,22 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None) # advance activations wp.launch( - _next_activation, - dim=(d.nworld, m.na), - inputs=[ - m.opt.timestep, - m.actuator_dyntype, - m.actuator_actlimited, - m.actuator_dynprm, - m.actuator_actrange, - d.act, - d.act_dot, - 1.0, - True, - ], - outputs=[d.act], + _next_activation, + dim=(d.nworld, m.nu), + inputs=[ + m.opt.timestep, + m.actuator_dyntype, + m.actuator_actadr, + m.actuator_actnum, + m.actuator_actlimited, + m.actuator_dynprm, + m.actuator_actrange, + d.act, + d.act_dot, + 1.0, + True, + ], + outputs=[d.act], ) wp.launch( @@ -285,7 +294,9 @@ def _euler_damp_qfrc_sparse( timestep = opt_timestep[worldid % opt_timestep.shape[0]] adr = dof_Madr[tid] - qM_integration_out[worldid, 0, adr] += timestep * dof_damping[worldid, tid] + qM_integration_out[worldid, 0, adr] += ( + timestep * dof_damping[worldid % dof_damping.shape[0], tid] + ) @cache_kernel @@ -379,10 +390,22 @@ def _rk_perturb_state( # activation if m.na and act_t0 is not None: wp.launch( - _next_activation, - dim=(d.nworld, m.na), - inputs=[m.opt.timestep, act_t0, d.act_dot, scale, False], - outputs=[d.act], + _next_activation, + dim=(d.nworld, m.nu), + inputs=[ + m.opt.timestep, + m.actuator_dyntype, + m.actuator_actadr, + m.actuator_actnum, + m.actuator_actlimited, + m.actuator_dynprm, + m.actuator_actrange, + act_t0, + d.act_dot, + scale, + False, + ], + outputs=[d.act], ) @@ -517,28 +540,35 @@ def fwd_position(m: Model, d: Data, factorize: bool = True): if m.opt.run_collision_detection: collision_driver.collision(m, d) constraint.make_constraint(m, d) + # TODO(team): remove False after island features are more complete + if False and not (m.opt.disableflags & DisableBit.ISLAND): + island.island(m, d) smooth.transmission(m, d) -# TODO(team): sparse actuator_moment version -@cache_kernel -def _actuator_velocity(nv: int): - @wp.kernel(module="unique", enable_backward=False) - def actuator_velocity( +@wp.kernel +def _actuator_velocity( # Data in: qvel_in: wp.array2d(dtype=float), - actuator_moment_in: wp.array3d(dtype=float), + moment_rownnz_in: wp.array2d(dtype=int), + moment_rowadr_in: wp.array2d(dtype=int), + moment_colind_in: wp.array2d(dtype=int), + actuator_moment_in: wp.array2d(dtype=float), # Data out: actuator_velocity_out: wp.array2d(dtype=float), - ): - worldid, actid = wp.tid() - moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(nv)) - qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv)) - moment_qvel_tile = wp.tile_map(wp.mul, moment_tile, qvel_tile) - actuator_velocity_tile = wp.tile_reduce(wp.add, moment_qvel_tile) - actuator_velocity_out[worldid, actid] = actuator_velocity_tile[0] +): + worldid, actid = wp.tid() - return actuator_velocity + rownnz = moment_rownnz_in[worldid, actid] + rowadr = moment_rowadr_in[worldid, actid] + + vel = float(0.0) + for i in range(rownnz): + sparseid = rowadr + i + colind = moment_colind_in[worldid, sparseid] + vel += actuator_moment_in[worldid, sparseid] * qvel_in[worldid, colind] + + actuator_velocity_out[worldid, actid] = vel @cache_kernel @@ -565,11 +595,17 @@ def _tendon_velocity(nv: int): def fwd_velocity(m: Model, d: Data): """Velocity-dependent computations.""" wp.launch_tiled( - _actuator_velocity(m.nv), - dim=(d.nworld, m.nu), - inputs=[d.qvel, d.actuator_moment], - outputs=[d.actuator_velocity], - block_dim=m.block_dim.actuator_velocity, + _actuator_velocity, + dim=(d.nworld, m.nu), + inputs=[ + d.qvel, + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + d.actuator_moment, + ], + outputs=[d.actuator_velocity], + block_dim=m.block_dim.actuator_velocity, ) # TODO(team): sparse version @@ -589,36 +625,36 @@ def fwd_velocity(m: Model, d: Data): @wp.kernel def _actuator_force( - # Model: - na: int, - opt_timestep: wp.array(dtype=float), - actuator_dyntype: wp.array(dtype=int), - actuator_gaintype: wp.array(dtype=int), - actuator_biastype: wp.array(dtype=int), - actuator_actadr: wp.array(dtype=int), - actuator_actnum: wp.array(dtype=int), - actuator_ctrllimited: wp.array(dtype=bool), - actuator_forcelimited: wp.array(dtype=bool), - actuator_actlimited: wp.array(dtype=bool), - actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_gainprm: wp.array2d(dtype=vec10f), - actuator_biasprm: wp.array2d(dtype=vec10f), - actuator_actearly: wp.array(dtype=bool), - actuator_ctrlrange: wp.array2d(dtype=wp.vec2), - actuator_forcerange: wp.array2d(dtype=wp.vec2), - actuator_actrange: wp.array2d(dtype=wp.vec2), - actuator_acc0: wp.array(dtype=float), - actuator_lengthrange: wp.array(dtype=wp.vec2), - # Data in: - act_in: wp.array2d(dtype=float), - ctrl_in: wp.array2d(dtype=float), - actuator_length_in: wp.array2d(dtype=float), - actuator_velocity_in: wp.array2d(dtype=float), - # In: - dsbl_clampctrl: int, - # Data out: - act_dot_out: wp.array2d(dtype=float), - actuator_force_out: wp.array2d(dtype=float), + # Model: + na: int, + opt_timestep: wp.array(dtype=float), + actuator_dyntype: wp.array(dtype=int), + actuator_gaintype: wp.array(dtype=int), + actuator_biastype: wp.array(dtype=int), + actuator_actadr: wp.array(dtype=int), + actuator_actnum: wp.array(dtype=int), + actuator_ctrllimited: wp.array(dtype=bool), + actuator_forcelimited: wp.array(dtype=bool), + actuator_actlimited: wp.array(dtype=bool), + actuator_dynprm: wp.array2d(dtype=vec10f), + actuator_gainprm: wp.array2d(dtype=vec10f), + actuator_biasprm: wp.array2d(dtype=vec10f), + actuator_actearly: wp.array(dtype=bool), + actuator_ctrlrange: wp.array2d(dtype=wp.vec2), + actuator_forcerange: wp.array2d(dtype=wp.vec2), + actuator_actrange: wp.array2d(dtype=wp.vec2), + actuator_acc0: wp.array2d(dtype=float), + actuator_lengthrange: wp.array2d(dtype=wp.vec2), + # Data in: + act_in: wp.array2d(dtype=float), + ctrl_in: wp.array2d(dtype=float), + actuator_length_in: wp.array2d(dtype=float), + actuator_velocity_in: wp.array2d(dtype=float), + # In: + dsbl_clampctrl: int, + # Data out: + act_dot_out: wp.array2d(dtype=float), + actuator_force_out: wp.array2d(dtype=float), ): worldid, uid = wp.tid() @@ -643,9 +679,11 @@ def _actuator_force( act = act_in[worldid, act_last] act_dot = (ctrl - act) / wp.max(dynprm[0], MJ_MINVAL) elif dyntype == DynType.MUSCLE: - dynprm = actuator_dynprm[worldid, uid] + dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], uid] act = act_in[worldid, act_last] act_dot = util_misc.muscle_dynamics(ctrl, act, dynprm) + elif dyntype == DynType.USER: + act_dot = 0.0 # set by act_dyn_callback else: # DynType.NONE act_dot = 0.0 @@ -681,20 +719,25 @@ def _actuator_force( elif gaintype == GainType.AFFINE: gain = gainprm[0] + gainprm[1] * length + gainprm[2] * velocity elif gaintype == GainType.MUSCLE: - acc0 = actuator_acc0[uid] - lengthrange = actuator_lengthrange[uid] + acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid] + lengthrange = actuator_lengthrange[ + worldid % actuator_lengthrange.shape[0], uid + ] gain = util_misc.muscle_gain(length, velocity, lengthrange, acc0, gainprm) + # GainType.USER: gain stays 0, modified by act_gain_callback # bias biastype = actuator_biastype[uid] biasprm = actuator_biasprm[worldid % actuator_biasprm.shape[0], uid] - bias = 0.0 # BiasType.NONE + bias = 0.0 # BiasType.NONE or BiasType.USER (modified by act_bias_callback) if biastype == BiasType.AFFINE: bias = biasprm[0] + biasprm[1] * length + biasprm[2] * velocity elif biastype == BiasType.MUSCLE: - acc0 = actuator_acc0[uid] - lengthrange = actuator_lengthrange[uid] + acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid] + lengthrange = actuator_lengthrange[ + worldid % actuator_lengthrange.shape[0], uid + ] bias = util_misc.muscle_bias(length, lengthrange, acc0, biasprm) force = gain * ctrl_act + bias @@ -752,32 +795,54 @@ def _tendon_actuator_force_clamp( @wp.kernel def _qfrc_actuator( - # Model: - nu: int, - ngravcomp: int, - jnt_actfrclimited: wp.array(dtype=bool), - jnt_actgravcomp: wp.array(dtype=int), - jnt_actfrcrange: wp.array2d(dtype=wp.vec2), - dof_jntid: wp.array(dtype=int), - # Data in: - actuator_moment_in: wp.array3d(dtype=float), - qfrc_gravcomp_in: wp.array2d(dtype=float), - actuator_force_in: wp.array2d(dtype=float), - # Data out: - qfrc_actuator_out: wp.array2d(dtype=float), + # Data in: + moment_rownnz_in: wp.array2d(dtype=int), + moment_rowadr_in: wp.array2d(dtype=int), + moment_colind_in: wp.array2d(dtype=int), + actuator_moment_in: wp.array2d(dtype=float), + actuator_force_in: wp.array2d(dtype=float), + # Data out: + qfrc_actuator_out: wp.array2d(dtype=float), +): + worldid, actid = wp.tid() + + rownnz = moment_rownnz_in[worldid, actid] + rowadr = moment_rowadr_in[worldid, actid] + + for i in range(rownnz): + sparseid = rowadr + i + colind = moment_colind_in[worldid, sparseid] + qfrc = ( + actuator_moment_in[worldid, sparseid] + * actuator_force_in[worldid, actid] + ) + wp.atomic_add(qfrc_actuator_out[worldid], colind, qfrc) + + +@wp.kernel +def _qfrc_actuator_gravcomp_limits( + # Model: + ngravcomp: int, + jnt_actfrclimited: wp.array(dtype=bool), + jnt_actgravcomp: wp.array(dtype=int), + jnt_actfrcrange: wp.array2d(dtype=wp.vec2), + dof_jntid: wp.array(dtype=int), + # Data in: + qfrc_gravcomp_in: wp.array2d(dtype=float), + qfrc_actuator_in: wp.array2d(dtype=float), + # Data out: + qfrc_actuator_out: wp.array2d(dtype=float), ): worldid, dofid = wp.tid() - - qfrc = float(0.0) - for uid in range(nu): - qfrc += actuator_moment_in[worldid, uid, dofid] * actuator_force_in[worldid, uid] - jntid = dof_jntid[dofid] + qfrc = qfrc_actuator_in[worldid, dofid] + # actuator-level gravity compensation, skip if added as passive force if ngravcomp and jnt_actgravcomp[jntid]: qfrc += qfrc_gravcomp_in[worldid, dofid] + # limits if jnt_actfrclimited[jntid]: frcrange = jnt_actfrcrange[worldid % jnt_actfrcrange.shape[0], jntid] qfrc = wp.clamp(qfrc, frcrange[0], frcrange[1]) @@ -825,6 +890,13 @@ def fwd_actuation(m: Model, d: Data): outputs=[d.act_dot, d.actuator_force], ) + if m.callback.act_dyn: + m.callback.act_dyn(m, d) + if m.callback.act_gain: + m.callback.act_gain(m, d) + if m.callback.act_bias: + m.callback.act_bias(m, d) + if m.ntendon: # total actuator force at tendon ten_actfrc = wp.zeros((d.nworld, m.ntendon), dtype=float) @@ -842,21 +914,33 @@ def fwd_actuation(m: Model, d: Data): outputs=[d.actuator_force], ) + # TODO(team): optimize performance + d.qfrc_actuator.zero_() wp.launch( - _qfrc_actuator, - dim=(d.nworld, m.nv), - inputs=[ - m.nu, - m.ngravcomp, - m.jnt_actfrclimited, - m.jnt_actgravcomp, - m.jnt_actfrcrange, - m.dof_jntid, - d.actuator_moment, - d.qfrc_gravcomp, - d.actuator_force, - ], - outputs=[d.qfrc_actuator], + _qfrc_actuator, + dim=(d.nworld, m.nu), + inputs=[ + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + d.actuator_moment, + d.actuator_force, + ], + outputs=[d.qfrc_actuator], + ) + wp.launch( + _qfrc_actuator_gravcomp_limits, + dim=(d.nworld, m.nv), + inputs=[ + m.ngravcomp, + m.jnt_actfrclimited, + m.jnt_actgravcomp, + m.jnt_actfrcrange, + m.dof_jntid, + d.qfrc_gravcomp, + d.qfrc_actuator, + ], + outputs=[d.qfrc_actuator], ) @@ -923,6 +1007,9 @@ def forward(m: Model, d: Data): if m.sensor_e_kinetic == 0: # not computed by sensor sensor.energy_vel(m, d) + if not (m.opt.disableflags & DisableBit.ACTUATION): + if m.callback.control: + m.callback.control(m, d) fwd_actuation(m, d) fwd_acceleration(m, d, factorize=True) @@ -971,6 +1058,10 @@ def step1(m: Model, d: Data): if m.sensor_e_kinetic == 0: # not computed by sensor sensor.energy_vel(m, d) + if not (m.opt.disableflags & DisableBit.ACTUATION): + if m.callback.control: + m.callback.control(m, d) + @event_scope def step2(m: Model, d: Data): 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 e6533dd6..e516739c 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -14,32 +14,24 @@ # ============================================================================== import dataclasses -import importlib.metadata -import warnings from typing import Any, Optional, Sequence +import warnings import mujoco -import numpy as np -import packaging.version -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import bvh +from mujoco.mjx.third_party.mujoco_warp._src import math as mjmath from mujoco.mjx.third_party.mujoco_warp._src import render_util from mujoco.mjx.third_party.mujoco_warp._src import smooth from mujoco.mjx.third_party.mujoco_warp._src import types from mujoco.mjx.third_party.mujoco_warp._src import warp_util - - -def _is_mujoco_dev() -> bool: - """Checks if mujoco version is > 3.4.0.""" - version_str = getattr(mujoco, "__version__", None) - if not version_str: - version_str = importlib.metadata.version("mujoco") - version_str = version_str.split("-")[0].split(".dev")[0] - return packaging.version.parse(version_str) > packaging.version.parse("3.4.0") - - -BLEEDING_EDGE_MUJOCO = _is_mujoco_dev() +from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL +from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN +from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType +from mujoco.mjx.third_party.mujoco_warp._src.types import vec10 +from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version +import numpy as np +import warp as wp def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array | None: @@ -218,6 +210,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.opt = opt m.stat = stat + m.callback = types.Callback() m.nv_pad = _get_padded_sizes( mjm.nv, 0, is_sparse(mjm), types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE @@ -585,17 +578,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.qM_mulm_madr.append(madr) m.qM_mulm_rowadr.append(len(m.qM_mulm_col)) - # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml - if BLEEDING_EDGE_MUJOCO: - m.flexedge_J_rownnz = mjm.flexedge_J_rownnz - m.flexedge_J_rowadr = mjm.flexedge_J_rowadr - m.flexedge_J_colind = mjm.flexedge_J_colind.reshape(-1) - else: - mjd = mujoco.MjData(mjm) - mujoco.mj_forward(mjm, mjd) - m.flexedge_J_rownnz = mjd.flexedge_J_rownnz - m.flexedge_J_rowadr = mjd.flexedge_J_rowadr - m.flexedge_J_colind = mjd.flexedge_J_colind.reshape(-1) + m.flexedge_J_rownnz = mjm.flexedge_J_rownnz + m.flexedge_J_rowadr = mjm.flexedge_J_rowadr + m.flexedge_J_colind = mjm.flexedge_J_colind.reshape(-1) # 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}) @@ -646,6 +631,16 @@ def _default_njmax(mjm: mujoco.MjModel, mjd: Optional[mujoco.MjData] = None) -> return int(valid_sizes[np.searchsorted(valid_sizes, njmax)]) +def _resolve_batch_size( + na: int | None, n: int | None, nworld: int, default: int +) -> int: + if na is not None: + return na + if n is not None: + return n * nworld + return default + + def make_data( mjm: mujoco.MjModel, nworld: int = 1, @@ -675,37 +670,36 @@ def make_data( if nconmax is None: nconmax = _default_nconmax(mjm) - if nconmax < 0: - raise ValueError("nconmax must be >= 0") - - if nccdmax is None: - nccdmax = nconmax - elif nccdmax < 0: - raise ValueError("nccdmax must be >= 0") - elif nccdmax > nconmax: - raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})") - if njmax is None: njmax = _default_njmax(mjm) + if nconmax < 0: + raise ValueError("nconmax must be >= 0") + if njmax < 0: raise ValueError("njmax must be >= 0") if nworld < 1: raise ValueError(f"nworld must be >= 1") - if naconmax is None: - naconmax = nworld * nconmax - elif naconmax < 0: + naconmax = _resolve_batch_size(naconmax, nconmax, nworld, 0) + if naconmax < 0: raise ValueError("naconmax must be >= 0") - if naccdmax is None: - naccdmax = nworld * nccdmax - elif naccdmax < 0: + naccdmax = _resolve_batch_size(naccdmax, nccdmax, nworld, naconmax) + if naccdmax < 0: raise ValueError("naccdmax must be >= 0") elif naccdmax > naconmax: raise ValueError(f"naccdmax ({naccdmax}) must be <= naconmax ({naconmax})") + if nccdmax is None: + nccdmax = nconmax + else: + if nccdmax < 0: + raise ValueError("nccdmax must be >= 0") + elif nccdmax > nconmax: + raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})") + sizes = dict({"*": 1}, **{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model) if f.type is int}) sizes["nmaxcondim"] = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max() sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1)) @@ -718,6 +712,17 @@ def make_data( contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)}) efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)}) + if SPARSE_CONSTRAINT_JACOBIAN: + efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int) + efc.J_rowadr = wp.zeros((nworld, njmax), dtype=int) + efc.J_colind = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=int) + efc.J = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=float) + else: + efc.J_rownnz = wp.zeros((nworld, 0), dtype=int) + efc.J_rowadr = wp.zeros((nworld, 0), dtype=int) + efc.J_colind = wp.zeros((nworld, 0, 0), dtype=int) + efc.J = wp.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=float) + # world body and static geom (attached to the world) poses are precomputed # this speeds up scenes with many static geoms (e.g. terrains) # TODO(team): remove this when we introduce dof islands + sleeping @@ -729,31 +734,65 @@ def make_data( mocap_id = mjm.body_mocapid[mocap_body] d_kwargs = { - "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float), - "contact": contact, - "efc": efc, - "nworld": nworld, - "naconmax": naconmax, - "naccdmax": naccdmax, - "njmax": njmax, - "qM": None, - "qLD": None, - # world body - "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat), - "xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33), - "ximat": wp.array(np.tile(mjd.ximat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33), - # static geoms - "geom_xpos": wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3), - "geom_xmat": wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33), - # mocap - "mocap_pos": wp.array(np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.vec3), - "mocap_quat": wp.array( - np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.quat - ), - # equality constraints - "eq_active": wp.array(np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), shape=(nworld, mjm.neq), dtype=bool), - # flexedge - "flexedge_J": None, + "qpos": wp.array( + np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float + ), + "contact": contact, + "efc": efc, + "nworld": nworld, + "naconmax": naconmax, + "naccdmax": naccdmax, + "njmax": njmax, + "njmax_pad": sizes["njmax_pad"], + "qM": None, + "qLD": None, + # world body + "xquat": wp.array( + np.tile(mjd.xquat, (nworld, 1)), + shape=(nworld, mjm.nbody), + dtype=wp.quat, + ), + "xmat": wp.array( + np.tile(mjd.xmat, (nworld, 1)), + shape=(nworld, mjm.nbody), + dtype=wp.mat33, + ), + "ximat": wp.array( + np.tile(mjd.ximat, (nworld, 1)), + shape=(nworld, mjm.nbody), + dtype=wp.mat33, + ), + # static geoms + "geom_xpos": wp.array( + np.tile(mjd.geom_xpos, (nworld, 1)), + shape=(nworld, mjm.ngeom), + dtype=wp.vec3, + ), + "geom_xmat": wp.array( + np.tile(mjd.geom_xmat, (nworld, 1)), + shape=(nworld, mjm.ngeom), + dtype=wp.mat33, + ), + # mocap + "mocap_pos": wp.array( + np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), + shape=(nworld, mjm.nmocap), + dtype=wp.vec3, + ), + "mocap_quat": wp.array( + np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)), + shape=(nworld, mjm.nmocap), + dtype=wp.quat, + ), + # equality constraints + "eq_active": wp.array( + np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), + shape=(nworld, mjm.neq), + dtype=bool, + ), + # island arrays + "nisland": None, + "tree_island": None, } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -769,7 +808,9 @@ def make_data( d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float) d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float) - d.flexedge_J = wp.zeros((nworld, 1, mjd.flexedge_J.size), dtype=float) + # island discovery arrays + d.nisland = wp.zeros((nworld,), dtype=int) + d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int) return d @@ -808,39 +849,43 @@ def put_data( if nconmax is None: nconmax = _default_nconmax(mjm, mjd) - if nconmax < 0: - raise ValueError("nconmax must be >= 0") - - if nccdmax is None: - nccdmax = nconmax - elif nccdmax < 0: - raise ValueError("nccdmax must be >= 0") - elif nccdmax > nconmax: - raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})") - if njmax is None: njmax = _default_njmax(mjm, mjd) + if nconmax < 0: + raise ValueError("nconmax must be >= 0") + if njmax < 0: raise ValueError("njmax must be >= 0") if nworld < 1: raise ValueError(f"nworld must be >= 1") - if naconmax is None: - if mjd.ncon > nconmax: - raise ValueError(f"nconmax overflow (nconmax must be >= {mjd.ncon})") - naconmax = nworld * nconmax + naconmax_is_input = naconmax is not None + naconmax = _resolve_batch_size(naconmax, nconmax, nworld, 0) + if naconmax < 0: + raise ValueError("naconmax must be >= 0") + + if not naconmax_is_input and mjd.ncon > nconmax: + raise ValueError(f"nconmax overflow (nconmax must be >= {mjd.ncon})") elif naconmax < mjd.ncon * nworld: raise ValueError(f"naconmax overflow (naconmax must be >= {mjd.ncon * nworld})") - if naccdmax is None: - naccdmax = nworld * nccdmax - elif naccdmax < 0: + naccdmax = _resolve_batch_size(naccdmax, nccdmax, nworld, naconmax) + + if naccdmax < 0: raise ValueError("naccdmax must be >= 0") elif naccdmax > naconmax: raise ValueError(f"naccdmax ({naccdmax}) must be <= naconmax ({naconmax})") + if nccdmax is None: + nccdmax = nconmax + else: + if nccdmax < 0: + raise ValueError("nccdmax must be >= 0") + elif nccdmax > nconmax: + raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})") + if mjd.nefc > njmax: raise ValueError(f"njmax overflow (njmax must be >= {mjd.nefc})") @@ -887,7 +932,7 @@ def put_data( contact.geomcollisionid = wp.empty((naconmax,), dtype=int) # TODO(team): set values # create efc - efc_kwargs = {"J": None} + efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None} for f in dataclasses.fields(types.Constraint): if f.name in efc_kwargs: @@ -900,31 +945,82 @@ def put_data( efc = types.Constraint(**efc_kwargs) - if mujoco.mj_isSparse(mjm): - efc_j = np.zeros((mjd.nefc, mjm.nv)) - mujoco.mju_sparse2dense(efc_j, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind) + if SPARSE_CONSTRAINT_JACOBIAN: + # TODO(team): process efc_J sparsity structure for nv row shift + efc.J_rownnz = wp.array( + np.full((nworld, njmax), mjm.nv, dtype=int), dtype=int + ) + efc.J_rowadr = wp.array( + np.tile( + np.arange(0, njmax * mjm.nv, mjm.nv) + if mjm.nv + else np.zeros(njmax, dtype=int), + (nworld, 1), + ), + dtype=int, + ) + efc.J_colind = wp.array( + np.tile(np.arange(mjm.nv), (nworld, njmax)).reshape((nworld, 1, -1)), + dtype=int, + ) + + mj_efc_J = np.zeros((mjd.nefc, mjm.nv)) + if mjd.nefc: + if mujoco.mj_isSparse(mjm): + mujoco.mju_sparse2dense( + mj_efc_J, + mjd.efc_J, + mjd.efc_J_rownnz, + mjd.efc_J_rowadr, + mjd.efc_J_colind, + ) + else: + mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv)) + efc_J = np.zeros((njmax, mjm.nv), dtype=float) + efc_J[: mjd.nefc, : mjm.nv] = mj_efc_J + efc.J = wp.array( + np.tile(efc_J.reshape(-1), (nworld, 1, 1)).reshape((nworld, 1, -1)), + dtype=float, + ) else: - efc_j = mjd.efc_J.reshape((mjd.nefc, mjm.nv)) - efc.J = np.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=f.type.dtype) - efc.J[:, : mjd.nefc, : mjm.nv] = np.tile(efc_j, (nworld, 1, 1)) - efc.J = wp.array(efc.J, dtype=float) + efc.J_rownnz = wp.zeros((nworld, 0), dtype=int) + efc.J_rowadr = wp.zeros((nworld, 0), dtype=int) + efc.J_colind = wp.zeros((nworld, 0, 0), dtype=int) + + mj_efc_J = np.zeros((mjd.nefc, mjm.nv)) + if mjd.nefc: + if mujoco.mj_isSparse(mjm): + mujoco.mju_sparse2dense( + mj_efc_J, + mjd.efc_J, + mjd.efc_J_rownnz, + mjd.efc_J_rowadr, + mjd.efc_J_colind, + ) + else: + mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv)) + efc_J = np.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=float) + efc_J[:, : mjd.nefc, : mjm.nv] = np.tile(mj_efc_J, (nworld, 1, 1)) + efc.J = wp.array(efc_J, dtype=float) # create data d_kwargs = { - "contact": contact, - "efc": efc, - "nworld": nworld, - "naconmax": naconmax, - "naccdmax": naccdmax, - "njmax": njmax, - # fields set after initialization: - "solver_niter": None, - "qM": None, - "qLD": None, - "ten_J": None, - "actuator_moment": None, - "flexedge_J": None, - "nacon": None, + "contact": contact, + "efc": efc, + "nworld": nworld, + "naconmax": naconmax, + "naccdmax": naccdmax, + "njmax": njmax, + "njmax_pad": sizes["njmax_pad"], + # fields set after initialization: + "solver_niter": None, + "qM": None, + "qLD": None, + "ten_J": None, + "nacon": None, + # island arrays + "nisland": None, + "tree_island": None, } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -950,19 +1046,32 @@ def put_data( d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float) d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float) - d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float) + # island arrays + d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int) + d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1)), dtype=int) - if mjm.ntendon: - ten_J = np.zeros((mjm.ntendon, mjm.nv)) - mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjm.ten_J_rownnz, mjm.ten_J_rowadr, mjm.ten_J_colind.reshape(-1)) - d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) + ten_J = np.zeros((mjm.ntendon, mjm.nv)) + if mujoco.mj_isSparse(mjm) or check_version("mujoco>=3.5.1.dev872479828"): + if mjm.ntendon: + if check_version("mujoco>=3.5.1.dev875093374"): + mujoco.mju_sparse2dense( + ten_J, + mjd.ten_J.reshape(-1), + mjm.ten_J_rownnz, + mjm.ten_J_rowadr, + mjm.ten_J_colind.reshape(-1), + ) + else: + mujoco.mju_sparse2dense( + ten_J, + mjd.ten_J.reshape(-1), + mjd.ten_J_rownnz, + mjd.ten_J_rowadr, + mjd.ten_J_colind.reshape(-1), + ) else: - d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), 0.0), dtype=float) - - # TODO(taylorhowell): sparse actuator_moment - actuator_moment = np.zeros((mjm.nu, mjm.nv)) - mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind) - d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float) + ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv)) + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) d.nacon = wp.array([mjd.ncon * nworld], dtype=int) @@ -1066,28 +1175,15 @@ def get_data_into( result.cinert[:] = d.cinert.numpy()[world_id] result.flexvert_xpos[:] = d.flexvert_xpos.numpy()[world_id] if mjm.nflexedge > 0: - # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml - if not BLEEDING_EDGE_MUJOCO: - m = put_model(mjm) - result.flexedge_J_rownnz[:] = m.flexedge_J_rownnz.numpy() - result.flexedge_J_rowadr[:] = m.flexedge_J_rowadr.numpy() - result.flexedge_J_colind[:, :] = m.flexedge_J_colind.numpy().reshape((mjm.nflexedge, mjm.nv)) - mujoco.mju_sparse2dense( - result.flexedge_J, - d.flexedge_J.numpy()[world_id].reshape(-1), - m.flexedge_J_rownnz.numpy(), - m.flexedge_J_rowadr.numpy(), - m.flexedge_J_colind.numpy(), - ) - else: - result.flexedge_J[:] = d.flexedge_J.numpy()[world_id].reshape(-1) + result.flexedge_J[:] = d.flexedge_J.numpy()[world_id].reshape(-1) result.flexedge_length[:] = d.flexedge_length.numpy()[world_id] result.flexedge_velocity[:] = d.flexedge_velocity.numpy()[world_id] result.actuator_length[:] = d.actuator_length.numpy()[world_id] - actuator_moment = d.actuator_moment.numpy()[world_id] - mujoco.mju_dense2sparse( - result.actuator_moment, actuator_moment, result.moment_rownnz, result.moment_rowadr, result.moment_colind - ) + result.moment_rownnz[:] = d.moment_rownnz.numpy()[world_id] + result.moment_rowadr[:] = d.moment_rowadr.numpy()[world_id] + if mjm.nu: + result.moment_colind[:] = d.moment_colind.numpy()[world_id] + result.actuator_moment[:] = d.actuator_moment.numpy()[world_id] result.crb[:] = d.crb.numpy()[world_id] result.qLDiagInv[:] = d.qLDiagInv.numpy()[world_id] result.ten_velocity[:] = d.ten_velocity.numpy()[world_id] @@ -1137,11 +1233,29 @@ def get_data_into( mujoco.mj_factorM(mjm, result) if nefc > 0: - if mujoco.mj_isSparse(mjm): - efc_J = d.efc.J.numpy()[world_id, efc_idx, : mjm.nv] - mujoco.mju_dense2sparse(result.efc_J, efc_J, result.efc_J_rownnz, result.efc_J_rowadr, result.efc_J_colind) + if SPARSE_CONSTRAINT_JACOBIAN: + efc_J = np.zeros((nefc, mjm.nv)) + mujoco.mju_sparse2dense( + efc_J, + d.efc.J.numpy()[world_id, 0], + d.efc.J_rownnz.numpy()[world_id, :nefc], + d.efc.J_rowadr.numpy()[world_id, :nefc], + d.efc.J_colind.numpy()[world_id, 0], + ) else: - result.efc_J[: nefc * mjm.nv] = d.efc.J.numpy()[world_id, :nefc, : mjm.nv].flatten() + efc_J = d.efc.J.numpy()[world_id, :nefc, : mjm.nv] + + # write to mujoco result (format depends on mj_isSparse) + if mujoco.mj_isSparse(mjm): + mujoco.mju_dense2sparse( + result.efc_J, + efc_J[efc_idx], + result.efc_J_rownnz, + result.efc_J_rowadr, + result.efc_J_colind, + ) + else: + result.efc_J[: nefc * mjm.nv] = efc_J[efc_idx].flatten() # efc result.efc_type[:] = d.efc.type.numpy()[world_id, efc_idx] @@ -1162,15 +1276,22 @@ def get_data_into( # tendon result.ten_length[:] = d.ten_length.numpy()[world_id] - # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml - if BLEEDING_EDGE_MUJOCO: + if check_version("mujoco>=3.5.1.dev869712136"): ten_J = d.ten_J.numpy()[world_id] + if check_version("mujoco>=3.5.1.dev875093374"): + ten_J_rownnz = mjm.ten_J_rownnz + ten_J_rowadr = mjm.ten_J_rowadr + ten_J_colind = mjm.ten_J_colind.reshape(-1) + else: + ten_J_rownnz = result.ten_J_rownnz + ten_J_rowadr = result.ten_J_rowadr + ten_J_colind = result.ten_J_colind.reshape(-1) mujoco.mju_dense2sparse( result.ten_J, ten_J, - mjm.ten_J_rownnz, - mjm.ten_J_rowadr, - mjm.ten_J_colind, + ten_J_rownnz, + ten_J_rowadr, + ten_J_colind, ) else: result.ten_J[:] = d.ten_J.numpy()[world_id] @@ -1182,6 +1303,12 @@ def get_data_into( # sensors result.sensordata[:] = d.sensordata.numpy()[world_id] + # islands + nisland = d.nisland.numpy()[world_id] + result.nisland = nisland + if nisland: + result.tree_island[:] = d.tree_island.numpy()[world_id] + def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): """Clear data, set defaults; optionally by world. @@ -1301,8 +1428,12 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): mocapid = body_mocapid[bodyid] if mocapid >= 0: - mocap_pos_out[worldid, mocapid] = body_pos[worldid, bodyid] - mocap_quat_out[worldid, mocapid] = body_quat[worldid, bodyid] + mocap_pos_out[worldid, mocapid] = body_pos[ + worldid % body_pos.shape[0], bodyid + ] + mocap_quat_out[worldid, mocapid] = body_quat[ + worldid % body_quat.shape[0], bodyid + ] @wp.kernel(module="unique", enable_backward=False) def reset_contact( @@ -1760,28 +1891,140 @@ def _compute_light_pos0( @wp.kernel def _copy_actuator_moment( - actid_target: int, - actuator_moment_in: wp.array3d(dtype=float), - act_moment_vec_out: wp.array2d(dtype=float), + 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), ): worldid = wp.tid() - nv = actuator_moment_in.shape[2] + nv = act_moment_vec_out.shape[1] for i in range(nv): - act_moment_vec_out[worldid, i] = actuator_moment_in[worldid, actid_target, i] + act_moment_vec_out[worldid, i] = 0.0 + rownnz = moment_rownnz_in[worldid, actid_target] + rowadr = moment_rowadr_in[worldid, actid_target] + for i in range(rownnz): + sparseid = rowadr + i + col = moment_colind_in[worldid, sparseid] + act_moment_vec_out[worldid, col] = actuator_moment_in[worldid, sparseid] @wp.kernel def _compute_actuator_acc0( - actid_target: int, - nv: int, - result_vec_in: wp.array2d(dtype=float), - actuator_acc0_out: wp.array(dtype=float), + actid_target: int, + nv: int, + result_vec_in: wp.array2d(dtype=float), + actuator_acc0_out: wp.array2d(dtype=float), ): worldid = wp.tid() norm_sq = float(0.0) for i in range(nv): norm_sq += result_vec_in[worldid, i] * result_vec_in[worldid, i] - actuator_acc0_out[actid_target] = wp.sqrt(norm_sq) + actuator_acc0_out[worldid, actid_target] = wp.sqrt(norm_sq) + + +@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), +): + worldid, dofid = wp.tid() + bodyid = dof_bodyid[dofid] + armature = dof_armature[worldid % dof_armature.shape[0], dofid] + buf = mjmath.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid]) + dof_M0_out[worldid, dofid] = armature + wp.dot(cdof_in[worldid, dofid], buf) + + +@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), + nv: int, + actuator_biasprm: wp.array2d(dtype=types.vec10f), +): + worldid, actid = wp.tid() + biastype = actuator_biastype[actid] + + # only affine bias (position actuators) + if biastype != BiasType.AFFINE: + return + + gainprm_id = worldid % actuator_gainprm.shape[0] + biasprm_id = worldid % actuator_biasprm.shape[0] + kp = actuator_gainprm[gainprm_id, actid][0] + + biasprm = actuator_biasprm[biasprm_id, actid] + # dampratio condition: gainprm[0] == -biasprm[1] and biasprm[2] > 0 + if wp.abs(kp + biasprm[1]) > MJ_MINVAL: + return + if biasprm[2] <= 0.0: + return + + dampratio = biasprm[2] + + # compute reflected mass: sum(dof_M0[j] / moment[i,j]^2) for active DOFs + mass = float(0.0) + rownnz = moment_rownnz_in[worldid, actid] + rowadr = moment_rowadr_in[worldid, actid] + for k in range(rownnz): + sparseid = rowadr + k + j = moment_colind_in[worldid, sparseid] + moment = actuator_moment_in[worldid, sparseid] + if wp.abs(moment) > MJ_MINVAL: + mass += dof_M0_in[worldid, j] / (moment * moment) + + damping = dampratio * 2.0 * wp.sqrt(kp * mass) + + # write -damping to biasprm[2] + new_biasprm = biasprm + new_biasprm[2] = -damping + actuator_biasprm[biasprm_id, actid] = new_biasprm + + +@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), + ntendon: int, + actuator_lengthrange_out: wp.array2d(dtype=wp.vec2), +): + worldid, actid = wp.tid() + trntype = actuator_trntype[actid] + id0 = actuator_trnid[actid][0] + gear0 = actuator_gear[worldid % actuator_gear.shape[0], actid][0] + + lr = wp.vec2(0.0, 0.0) + + if trntype == TrnType.JOINT or trntype == TrnType.JOINTINPARENT: + if jnt_limited[id0]: + rng = jnt_range[worldid % jnt_range.shape[0], id0] + if gear0 > 0.0: + lr = wp.vec2(rng[0] * gear0, rng[1] * gear0) + else: + lr = wp.vec2(rng[1] * gear0, rng[0] * gear0) + elif trntype == TrnType.TENDON: + if ntendon > 0 and tendon_limited[id0]: + rng = tendon_range[worldid % tendon_range.shape[0], id0] + if gear0 > 0.0: + lr = wp.vec2(rng[0] * gear0, rng[1] * gear0) + else: + lr = wp.vec2(rng[1] * gear0, rng[0] * gear0) + + actuator_lengthrange_out[worldid, actid] = lr # kernel_analyzer: on @@ -1823,6 +2066,9 @@ def set_const_0(m: types.Model, d: types.Data): - cam_pos0, cam_poscom0, cam_mat0: camera references - light_pos0, light_poscom0, light_dir0: light references - actuator_acc0: acceleration from unit actuator force + - actuator_biasprm[2] (dampratio resolution): for position actuators where + gainprm[0] == -biasprm[1] and biasprm[2] > 0, converts dampratio to + damping via biasprm[2] = -dampratio * 2 * sqrt(kp * reflected_mass) Args: m: The model containing kinematic and dynamic information (device). @@ -1954,10 +2200,46 @@ def set_const_0(m: types.Model, d: types.Data): act_result_vec = wp.zeros((d.nworld, m.nv), dtype=float) for actid in range(m.nu): - wp.launch(_copy_actuator_moment, dim=d.nworld, inputs=[actid, d.actuator_moment], outputs=[act_moment_vec]) + wp.launch( + _copy_actuator_moment, + dim=d.nworld, + inputs=[ + actid, + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + d.actuator_moment, + ], + outputs=[act_moment_vec], + ) smooth.solve_m(m, d, act_result_vec, act_moment_vec) wp.launch(_compute_actuator_acc0, dim=d.nworld, inputs=[actid, m.nv, act_result_vec], outputs=[m.actuator_acc0]) + # resolve dampratio: compute dof_M0, then convert dampratio to damping + if m.nu > 0 and m.nv > 0: + dof_M0 = wp.zeros((d.nworld, m.nv), dtype=float) + wp.launch( + _compute_dof_M0, + dim=(d.nworld, m.nv), + inputs=[m.dof_bodyid, m.dof_armature, d.cdof, d.crb], + outputs=[dof_M0], + ) + wp.launch( + _resolve_dampratio, + dim=(d.nworld, m.nu), + inputs=[ + m.actuator_biastype, + m.actuator_gainprm, + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + d.actuator_moment, + dof_M0, + m.nv, + ], + outputs=[m.actuator_biasprm], + ) + wp.copy(d.qpos, qpos_saved) @@ -1973,12 +2255,16 @@ def set_const(m: types.Model, d: types.Data): Field | Notes ---------------------------------|---------------------------------------------- qpos0, qpos_spring | - body_mass, body_inertia, | Mass and inertia are usually scaled together + body_mass, body_inertia, | Mass and inertia are usually scaled + together body_ipos, body_iquat | since inertia is sum(m * r^2). - body_pos, body_quat | Unsafe for static bodies (invalidates BVH). - body_gravcomp | If changing from 0 to >0 bodies, required. + body_pos, body_quat | Unsafe for static bodies (invalidates + BVH). + body_gravcomp | If changing from 0 to >0 bodies, + required. dof_armature | - eq_data | For connect/weld, offsets computed if not set. + eq_data | For connect/weld, offsets computed if not + set. hfield_size | tendon_stiffness, tendon_damping | Only if changing from/to zero. actuator_gainprm, actuator_biasprm | For position actuators with dampratio. @@ -2004,8 +2290,9 @@ def set_const(m: types.Model, d: types.Data): - cam_pos0, cam_poscom0, cam_mat0: camera references - light_pos0, light_poscom0, light_dir0: light references - actuator_acc0: acceleration from unit actuator force + - actuator_biasprm[2] (dampratio resolution) - Skips: dof_M0, actuator_length0 (not in mjwarp). + Skips: actuator_length0 (not in mjwarp). Args: m: The model containing kinematic and dynamic information (device). @@ -2015,6 +2302,39 @@ def set_const(m: types.Model, d: types.Data): set_const_0(m, d) +def set_length_range(m: types.Model, d: types.Data, index: int = -1): + """Compute feasible actuator length ranges from joint/tendon limits. + + For joint and tendon transmissions with limits, copies the range directly + from jnt_range or tendon_range scaled by gear. Actuators without limits + keep (0, 0). This covers the common robotics use case; simulation-based + computation for general transmissions is not yet implemented. + + Args: + m: The model containing kinematic and dynamic information (device). + d: The data object (unused, kept for API compatibility with MuJoCo C). + index: Actuator index to compute for, or -1 for all actuators. + """ + if m.nu == 0: + return + + wp.launch( + _set_length_range, + dim=(d.nworld, m.nu), + inputs=[ + m.actuator_trntype, + m.actuator_trnid, + m.actuator_gear, + m.jnt_limited, + m.jnt_range, + m.tendon_limited, + m.tendon_range, + m.ntendon, + ], + outputs=[m.actuator_lengthrange], + ) + + def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any] | Sequence[str]): """Overrides model parameters. @@ -2321,26 +2641,25 @@ def create_render_context( cam_res_arr = wp.array(active_cam_res, dtype=wp.vec2i) - if render_rgb and isinstance(render_rgb, bool): + if render_rgb is None: + render_rgb = [ + mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB + for i in active_cam_indices + ] + elif isinstance(render_rgb, bool): render_rgb = [render_rgb] * ncam - elif render_rgb is None: - # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml - if BLEEDING_EDGE_MUJOCO: - render_rgb = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB for i in active_cam_indices] - else: - render_rgb = [True] * ncam - if render_depth and isinstance(render_depth, bool): + if render_depth is None: + render_depth = [ + mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH + for i in active_cam_indices + ] + if isinstance(render_depth, bool): render_depth = [render_depth] * ncam - elif render_depth is None: - # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml - if BLEEDING_EDGE_MUJOCO: - render_depth = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH for i in active_cam_indices] - else: - render_depth = [True] * ncam assert len(render_rgb) == ncam and len(render_depth) == ncam, ( - f"Render RGB and depth must be provided for all active cameras (got {len(render_rgb)}, {len(render_depth)}, expected {ncam})" + "render_rgb and render_depth must be a bool or a list of bools with" + f" length {ncam}" ) rgb_adr = -1 * np.ones(ncam, dtype=int) @@ -2364,10 +2683,7 @@ def create_render_context( ray = wp.zeros(int(total), dtype=wp.vec3) - # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml - cam_projection = np.zeros(mjm.ncam, dtype=int) - if BLEEDING_EDGE_MUJOCO: - cam_projection = mjm.cam_projection + cam_projection = mjm.cam_projection offset = 0 for idx, cam_id in enumerate(active_cam_indices): 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 1b93db0c..b5ae2846 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py @@ -13,12 +13,12 @@ # limitations under the License. # ============================================================================== -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import types from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType from mujoco.mjx.third_party.mujoco_warp._src.types import EqType from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp @wp.kernel @@ -44,7 +44,7 @@ def _tree_edges( # Out: tree_tree: wp.array3d(dtype=int), # kernel_analyzer: off ): - """Find tree edges. Launch: (nworld, njmax).""" + """Find tree edges.""" worldid, efcid = wp.tid() # skip if beyond active constraints @@ -176,3 +176,89 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)): ], outputs=[tree_tree], ) + + +@wp.kernel +def _flood_fill( + # Model: + ntree: int, + # In: + tree_tree_in: wp.array3d(dtype=int), + labels_in: wp.array2d(dtype=int), + stack_in: wp.array2d(dtype=int), + # Data out: + nisland_out: wp.array(dtype=int), + tree_island_out: wp.array2d(dtype=int), + # Out: + stack_out: wp.array2d(dtype=int), +): + """DFS flood fill to discover islands using tree_tree matrix.""" + worldid = wp.tid() + nisland = int(0) + + # iterate over trees + for i in range(ntree): + # already assigned + if labels_in[worldid, i] != -1: + continue + + # check if tree has any edges + has_edge = int(0) + for j in range(ntree): + if tree_tree_in[worldid, i, j] != 0: + has_edge = 1 + break + if has_edge == 0: + continue + + # DFS: push i onto stack + nstack = int(0) + stack_out[worldid, nstack] = i + nstack = nstack + 1 + + while nstack > 0: + # pop v from stack + nstack = nstack - 1 + v = stack_in[worldid, nstack] + + # already assigned + if labels_in[worldid, v] != -1: + continue + + # assign to current island + tree_island_out[worldid, v] = nisland + + # push neighbors + for neighbor in range(ntree): + if tree_tree_in[worldid, v, neighbor] != 0: + if labels_in[worldid, neighbor] == -1: + stack_out[worldid, nstack] = neighbor + nstack = nstack + 1 + + # island filled + nisland = nisland + 1 + + nisland_out[worldid] = nisland + + +@event_scope +def island(m: types.Model, d: types.Data): + """Discover constraint islands.""" + if m.ntree == 0: + d.nisland.zero_() + return + + # Step 1: Find tree edges + tree_tree = wp.zeros((d.nworld, m.ntree, m.ntree), dtype=int) + tree_edges(m, d, tree_tree) + + # Step 2: DFS flood fill + d.tree_island.fill_(-1) + stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int) + + wp.launch( + _flood_fill, + dim=d.nworld, + inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch], + outputs=[d.nisland, d.tree_island, stack_scratch], + ) 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 77ed86c0..0bcd13af 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -265,28 +265,28 @@ def _gravity_force( @wp.kernel def _fluid_force( - # Model: - opt_density: wp.array(dtype=float), - opt_viscosity: wp.array(dtype=float), - opt_wind: wp.array(dtype=wp.vec3), - body_rootid: wp.array(dtype=int), - body_geomnum: wp.array(dtype=int), - body_geomadr: wp.array(dtype=int), - body_mass: wp.array2d(dtype=float), - body_inertia: wp.array2d(dtype=wp.vec3), - geom_type: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - geom_fluid: wp.array2d(dtype=float), - body_fluid_ellipsoid: wp.array(dtype=bool), - # Data in: - xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - # Out: - fluid_applied_out: wp.array2d(dtype=wp.spatial_vector), + # Model: + opt_wind: wp.array(dtype=wp.vec3), + opt_density: wp.array(dtype=float), + opt_viscosity: wp.array(dtype=float), + body_rootid: wp.array(dtype=int), + body_geomnum: wp.array(dtype=int), + body_geomadr: wp.array(dtype=int), + body_mass: wp.array2d(dtype=float), + body_inertia: wp.array2d(dtype=wp.vec3), + geom_type: wp.array(dtype=int), + geom_size: wp.array2d(dtype=wp.vec3), + geom_fluid: wp.array2d(dtype=float), + body_fluid_ellipsoid: wp.array(dtype=bool), + # Data in: + xipos_in: wp.array2d(dtype=wp.vec3), + ximat_in: wp.array2d(dtype=wp.mat33), + geom_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xmat_in: wp.array2d(dtype=wp.mat33), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cvel_in: wp.array2d(dtype=wp.spatial_vector), + # Out: + fluid_applied_out: wp.array2d(dtype=wp.spatial_vector), ): """Computes body-space fluid forces for both inertia-box and ellipsoid models.""" worldid, bodyid = wp.tid() @@ -495,29 +495,29 @@ def _fluid(m: Model, d: Data): fluid_applied = wp.empty((d.nworld, m.nbody), dtype=wp.spatial_vector) wp.launch( - _fluid_force, - dim=(d.nworld, m.nbody), - inputs=[ - m.opt.density, - m.opt.viscosity, - m.opt.wind, - m.body_rootid, - m.body_geomnum, - m.body_geomadr, - m.body_mass, - m.body_inertia, - m.geom_type, - m.geom_size, - m.geom_fluid, - m.body_fluid_ellipsoid, - d.xipos, - d.ximat, - d.geom_xpos, - d.geom_xmat, - d.subtree_com, - d.cvel, - ], - outputs=[fluid_applied], + _fluid_force, + dim=(d.nworld, m.nbody), + inputs=[ + m.opt.wind, + m.opt.density, + m.opt.viscosity, + m.body_rootid, + m.body_geomnum, + m.body_geomadr, + m.body_mass, + m.body_inertia, + m.geom_type, + m.geom_size, + m.geom_fluid, + m.body_fluid_ellipsoid, + d.xipos, + d.ximat, + d.geom_xpos, + d.geom_xmat, + d.subtree_com, + d.cvel, + ], + outputs=[fluid_applied], ) support.apply_ft(m, d, fluid_applied, d.qfrc_fluid, False) @@ -846,3 +846,6 @@ def passive(m: Model, d: Data): d.qfrc_passive, ], ) + + if m.callback.passive: + m.callback.passive(m, d) 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 befc9820..5eaea821 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py @@ -66,19 +66,19 @@ def _ray_eliminate( ) -> bool: """Eliminate ray.""" bodyid = geom_bodyid[geomid] - matid = geom_matid[geomid] + matid = geom_matid[geomid] # kernel_analyzer: ignore # body exclusion if bodyid == bodyexclude: return True # invisible geom exclusion - if matid < 0 and geom_rgba[geomid][3] == 0.0: + if matid < 0 and geom_rgba[geomid][3] == 0.0: # kernel_analyzer: ignore return True # invisible material exclusion if matid >= 0: - if mat_rgba[matid][3] == 0.0: + if mat_rgba[matid][3] == 0.0: # kernel_analyzer: ignore return True # static exclusion 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 79b79c06..28a4284f 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py @@ -552,7 +552,13 @@ def render(m: Model, d: Data, rc: RenderContext): return if render_depth[cam_idx]: - depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist + # Planar depth: project Euclidean distance onto the camera's optical axis. + # In camera-local coordinates, the optical axis is -Z. The Z-component of the + # normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ) + # between the ray and the optical axis. + depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist * ( + -ray_dir_local_cam[2] + ) if not render_rgb[cam_idx]: return 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 a152f468..859ddf8a 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py @@ -15,16 +15,12 @@ from typing import Any, Tuple -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import ray from mujoco.mjx.third_party.mujoco_warp._src import smooth from mujoco.mjx.third_party.mujoco_warp._src import support from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import get_sdf_params from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType @@ -32,9 +28,13 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import DataType from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import JointType +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType +from mujoco.mjx.third_party.mujoco_warp._src.types import Stage from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6 @@ -43,6 +43,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -898,6 +899,9 @@ def sensor_pos(m: Model, d: Data): ], ) + if m.callback.sensor: + m.callback.sensor(m, d, Stage.POS) + @wp.func def _velocimeter( @@ -1437,6 +1441,9 @@ def sensor_vel(m: Model, d: Data): ], ) + if m.callback.sensor: + m.callback.sensor(m, d, Stage.VEL) + @wp.func def _accelerometer( @@ -2073,67 +2080,92 @@ def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3: @wp.kernel -def _sensor_tactile( - # Model: - body_rootid: wp.array(dtype=int), - body_weldid: wp.array(dtype=int), - oct_child: wp.array(dtype=vec8i), - oct_aabb: wp.array2d(dtype=wp.vec3), - oct_coeff: wp.array(dtype=vec8), - geom_type: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - geom_size: wp.array2d(dtype=wp.vec3), - mesh_vertadr: wp.array(dtype=int), - mesh_normaladr: wp.array(dtype=int), - mesh_vert: wp.array(dtype=wp.vec3), - mesh_normal: wp.array(dtype=wp.vec3), - mesh_quat: wp.array(dtype=wp.quat), - sensor_objid: wp.array(dtype=int), - sensor_refid: wp.array(dtype=int), - sensor_dim: wp.array(dtype=int), - sensor_adr: wp.array(dtype=int), - plugin: wp.array(dtype=int), - plugin_attr: wp.array(dtype=wp.vec3f), - geom_plugin_index: wp.array(dtype=int), - taxel_vertadr: wp.array(dtype=int), - taxel_sensorid: wp.array(dtype=int), - # Data in: - geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cvel_in: wp.array2d(dtype=wp.spatial_vector), - contact_geom_in: wp.array(dtype=wp.vec2i), - contact_worldid_in: wp.array(dtype=int), - nacon_in: wp.array(dtype=int), - # Data out: - sensordata_out: wp.array2d(dtype=float), +def _preprocess_tactile_contacts( + # Model: + body_weldid: wp.array(dtype=int), + geom_bodyid: wp.array(dtype=int), + # Data in: + contact_geom_in: wp.array(dtype=wp.vec2i), + contact_worldid_in: wp.array(dtype=int), + nacon_in: wp.array(dtype=int), + # Out: + weld_geom_count_out: wp.array2d(dtype=int), + weld_geom_list_out: wp.array3d(dtype=int), ): - conid, taxelid = wp.tid() - - if conid >= nacon_in[0]: + conid = wp.tid() + ncon = nacon_in[0] + if conid >= ncon: return - worldid = contact_worldid_in[conid] + contact_geom = contact_geom_in[conid] + weld1 = body_weldid[geom_bodyid[contact_geom[0]]] + weld2 = body_weldid[geom_bodyid[contact_geom[1]]] + geom1 = contact_geom[0] + geom2 = contact_geom[1] + + for side in range(2): + if side == 0: + weld = weld1 + geom = geom2 + else: + weld = weld2 + geom = geom1 + + idx = wp.atomic_add(weld_geom_count_out[worldid], weld, 1) + if idx < MJ_MAXCONPAIR: + weld_geom_list_out[worldid, weld, idx] = geom + + +@wp.kernel +def _sensor_tactile( + # Model: + body_rootid: wp.array(dtype=int), + body_weldid: wp.array(dtype=int), + oct_child: wp.array(dtype=vec8i), + oct_aabb: wp.array2d(dtype=wp.vec3), + oct_coeff: wp.array(dtype=vec8), + geom_type: wp.array(dtype=int), + geom_bodyid: wp.array(dtype=int), + geom_size: wp.array2d(dtype=wp.vec3), + mesh_vertadr: wp.array(dtype=int), + mesh_vertnum: wp.array(dtype=int), + mesh_octadr: wp.array(dtype=int), + mesh_normaladr: wp.array(dtype=int), + mesh_normalnum: wp.array(dtype=int), + mesh_vert: wp.array(dtype=wp.vec3), + mesh_normal: wp.array(dtype=wp.vec3), + mesh_quat: wp.array(dtype=wp.quat), + sensor_objid: wp.array(dtype=int), + sensor_refid: wp.array(dtype=int), + sensor_dim: wp.array(dtype=int), + sensor_adr: wp.array(dtype=int), + plugin: wp.array(dtype=int), + plugin_attr: wp.array(dtype=wp.vec3f), + geom_plugin_index: wp.array(dtype=int), + taxel_vertadr: wp.array(dtype=int), + taxel_sensorid: wp.array(dtype=int), + # Data in: + geom_xpos_in: wp.array2d(dtype=wp.vec3), + geom_xmat_in: wp.array2d(dtype=wp.mat33), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cvel_in: wp.array2d(dtype=wp.spatial_vector), + # In: + weld_geom_count_in: wp.array2d(dtype=int), + weld_geom_list_in: wp.array3d(dtype=int), + # Data out: + sensordata_out: wp.array2d(dtype=float), +): + worldid, taxelid = wp.tid() - # get sensor_id sensor_id = taxel_sensorid[taxelid] - - # get parent weld id mesh_id = sensor_objid[sensor_id] geom_id = sensor_refid[sensor_id] parent_body = geom_bodyid[geom_id] parent_weld = body_weldid[parent_body] - # contact geom - body1 = body_weldid[geom_bodyid[contact_geom_in[conid][0]]] - body2 = body_weldid[geom_bodyid[contact_geom_in[conid][1]]] - if body1 == parent_weld: - geom = contact_geom_in[conid][1] - elif body2 == parent_weld: - geom = contact_geom_in[conid][0] - else: + geom_count = weld_geom_count_in[worldid, parent_weld] + if geom_count == 0: return - body = geom_bodyid[geom] # vertex local position vertid = taxel_vertadr[taxelid] - mesh_vertadr[mesh_id] @@ -2143,57 +2175,108 @@ def _sensor_tactile( xpos = geom_xmat_in[worldid, geom_id] @ pos xpos += geom_xpos_in[worldid, geom_id] - # position in other geom frame - tmp = xpos - geom_xpos_in[worldid, geom] - lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp + has_frame = mesh_normalnum[mesh_id] == 3 * mesh_vertnum[mesh_id] + normal_stride = 3 if has_frame else 1 + offset = mesh_normaladr[mesh_id] + normal_stride * vertid + quat = mesh_quat[mesh_id] + normal = math.rot_vec_quat(mesh_normal[offset], quat) + tang1 = wp.vec3(0.0, 0.0, 0.0) + tang2 = wp.vec3(0.0, 0.0, 0.0) + if has_frame: + tang1 = math.rot_vec_quat(mesh_normal[offset + 1], quat) + tang2 = math.rot_vec_quat(mesh_normal[offset + 2], quat) - plugin_id = geom_plugin_index[geom] + for g in range(MJ_MAXCONPAIR): + if g >= geom_count: + break - contact_type = geom_type[geom] + geom = weld_geom_list_in[worldid, parent_weld, g] + if geom < 0: + continue - plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params( - oct_child, - oct_aabb, - oct_coeff, - plugin, - plugin_attr, - contact_type, - geom_size[worldid % geom_size.shape[0], geom], - plugin_id, - mesh_id, - ) + is_dup = int(0) + for j in range(g): + if weld_geom_list_in[worldid, parent_weld, j] == geom: + is_dup = int(1) + break + if is_dup == int(1): + continue - depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0) - if depth >= 0.0: - return + body = geom_bodyid[geom] - # get velocity in global - vel_sensor = _transform_spatial(cvel_in[worldid, parent_weld], xpos - subtree_com_in[worldid, body_rootid[parent_weld]]) - vel_other = _transform_spatial( - cvel_in[worldid, body], geom_xpos_in[worldid, geom] - subtree_com_in[worldid, body_rootid[body]] - ) - vel_rel = vel_sensor - vel_other + tmp = xpos - geom_xpos_in[worldid, geom] + lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp - # get contact force/torque, rotate into node frame - offset = mesh_normaladr[mesh_id] + 3 * vertid - normal = math.rot_vec_quat(mesh_normal[offset], mesh_quat[mesh_id]) - tang1 = math.rot_vec_quat(mesh_normal[offset + 1], mesh_quat[mesh_id]) - tang2 = math.rot_vec_quat(mesh_normal[offset + 2], mesh_quat[mesh_id]) - kMaxDepth = 0.05 - pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL) - force = wp.mul(normal, pressure) + plugin_id = geom_plugin_index[geom] + contact_type = geom_type[geom] - # one row of mat^T * force - forceT = wp.vec3() - forceT[0] = wp.dot(force, normal) - forceT[1] = wp.abs(wp.dot(vel_rel, tang1)) - forceT[2] = wp.abs(wp.dot(vel_rel, tang2)) + plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params( + oct_child, + oct_aabb, + oct_coeff, + mesh_octadr, + plugin, + plugin_attr, + contact_type, + geom_size[worldid % geom_size.shape[0], geom], + plugin_id, + mesh_id, + ) - # add to sensor output - dim = sensor_dim[sensor_id] / 3 - wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0]) - wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1]) - wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2]) + depth = wp.min( + sdf( + contact_type, + lpos, + plugin_attributes, + plugin_index, + volume_data, + mesh_data, + ), + 0.0, + ) + if depth >= 0.0: + continue + + vel_sensor = _transform_spatial( + cvel_in[worldid, parent_weld], + xpos - subtree_com_in[worldid, body_rootid[parent_weld]], + ) + vel_other = _transform_spatial( + cvel_in[worldid, body], + geom_xpos_in[worldid, geom] + - subtree_com_in[worldid, body_rootid[body]], + ) + vel_rel = vel_sensor - vel_other + + kMaxDepth = 0.05 + pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL) + force = wp.mul(normal, pressure) + + forceT = wp.vec3(0.0, 0.0, 0.0) + forceT[0] = wp.dot(force, normal) + if has_frame: + forceT[1] = wp.abs(wp.dot(vel_rel, tang1)) + forceT[2] = wp.abs(wp.dot(vel_rel, tang2)) + + dim = sensor_dim[sensor_id] // 3 + wp.atomic_add( + sensordata_out, + worldid, + sensor_adr[sensor_id] + 0 * dim + vertid, + forceT[0], + ) + wp.atomic_add( + sensordata_out, + worldid, + sensor_adr[sensor_id] + 1 * dim + vertid, + forceT[1], + ) + wp.atomic_add( + sensordata_out, + worldid, + sensor_adr[sensor_id] + 2 * dim + vertid, + forceT[2], + ) @wp.func @@ -2421,43 +2504,63 @@ def sensor_acc(m: Model, d: Data): ], ) + weld_geom_count = wp.zeros((d.nworld, m.nbody), dtype=int) + weld_geom_list = wp.full((d.nworld, m.nbody, MJ_MAXCONPAIR), -1, dtype=int) wp.launch( - _sensor_tactile, - dim=(d.naconmax, m.nsensortaxel), - inputs=[ - m.body_rootid, - m.body_weldid, - m.oct_child, - m.oct_aabb, - m.oct_coeff, - m.geom_type, - m.geom_bodyid, - m.geom_size, - m.mesh_vertadr, - m.mesh_normaladr, - m.mesh_vert, - m.mesh_normal, - m.mesh_quat, - m.sensor_objid, - m.sensor_refid, - m.sensor_dim, - m.sensor_adr, - m.plugin, - m.plugin_attr, - m.geom_plugin_index, - m.taxel_vertadr, - m.taxel_sensorid, - d.geom_xpos, - d.geom_xmat, - d.subtree_com, - d.cvel, - d.contact.geom, - d.contact.worldid, - d.nacon, - ], - outputs=[ - d.sensordata, - ], + _preprocess_tactile_contacts, + dim=d.naconmax, + inputs=[ + m.body_weldid, + m.geom_bodyid, + d.contact.geom, + d.contact.worldid, + d.nacon, + ], + outputs=[ + weld_geom_count, + weld_geom_list, + ], + ) + + wp.launch( + _sensor_tactile, + dim=(d.nworld, m.nsensortaxel), + inputs=[ + m.body_rootid, + m.body_weldid, + m.oct_child, + m.oct_aabb, + m.oct_coeff, + m.geom_type, + m.geom_bodyid, + m.geom_size, + m.mesh_vertadr, + m.mesh_vertnum, + m.mesh_octadr, + m.mesh_normaladr, + m.mesh_normalnum, + m.mesh_vert, + m.mesh_normal, + m.mesh_quat, + m.sensor_objid, + m.sensor_refid, + m.sensor_dim, + m.sensor_adr, + m.plugin, + m.plugin_attr, + m.geom_plugin_index, + m.taxel_vertadr, + m.taxel_sensorid, + d.geom_xpos, + d.geom_xmat, + d.subtree_com, + d.cvel, + weld_geom_count, + weld_geom_list, + ], + outputs=[ + d.sensordata, + ], ) sensor_contact_nmatch = wp.empty((d.nworld, m.nsensorcontact), dtype=int) @@ -2616,6 +2719,9 @@ def sensor_acc(m: Model, d: Data): ], ) + if m.callback.sensor: + m.callback.sensor(m, d, Stage.ACC) + @wp.kernel def _energy_pos_zero( 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 ff5ddbe7..51b36403 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py @@ -14,29 +14,29 @@ # ============================================================================== -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import support from mujoco.mjx.third_party.mujoco_warp._src import util_misc -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL -from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import CamLightType from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType from mujoco.mjx.third_party.mujoco_warp._src.types import Data from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit from mujoco.mjx.third_party.mujoco_warp._src.types import EqType from mujoco.mjx.third_party.mujoco_warp._src.types import JointType +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL +from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType +from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType -from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType -from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10 from mujoco.mjx.third_party.mujoco_warp._src.types import vec11 +from mujoco.mjx.third_party.mujoco_warp._src.types import vec5 +from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -239,28 +239,26 @@ def _flex_vertices( @wp.kernel def _flex_edges( - # Model: - nflex: int, - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_dofadr: wp.array(dtype=int), - dof_bodyid: 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), - # 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), - # Data out: - flexedge_J_out: wp.array3d(dtype=float), - flexedge_length_out: wp.array2d(dtype=float), - flexedge_velocity_out: wp.array2d(dtype=float), + # Model: + nflex: int, + body_rootid: 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), + # 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), + # Data out: + flexedge_J_out: wp.array2d(dtype=float), + flexedge_length_out: wp.array2d(dtype=float), + flexedge_velocity_out: wp.array2d(dtype=float), ): worldid, edgeid = wp.tid() for i in range(nflex): @@ -285,64 +283,40 @@ def _flex_edges( dofi = body_dofadr[b1] dofj = body_dofadr[b2] - dofi0 = dofi + 0 - dofi1 = dofi + 1 - dofi2 = dofi + 2 - dofj0 = dofj + 0 - dofj1 = dofj + 1 - dofj2 = dofj + 2 - vel1 = wp.vec3(qvel_in[worldid, dofi0], qvel_in[worldid, dofi1], qvel_in[worldid, dofi2]) - vel2 = wp.vec3(qvel_in[worldid, dofj0], qvel_in[worldid, dofj1], qvel_in[worldid, dofj2]) + vel1 = wp.vec3( + qvel_in[worldid, dofi], + qvel_in[worldid, dofi + 1], + qvel_in[worldid, dofi + 2], + ) + vel2 = wp.vec3( + qvel_in[worldid, dofj], + qvel_in[worldid, dofj + 1], + qvel_in[worldid, dofj + 2], + ) flexedge_velocity_out[worldid, edgeid] = wp.dot(vel2 - vel1, edge) rowadr = flexedge_J_rowadr[edgeid] - sparseid0 = rowadr + 0 - sparseid1 = rowadr + 1 - sparseid2 = rowadr + 2 - sparseid3 = rowadr + 3 - sparseid4 = rowadr + 4 - sparseid5 = rowadr + 5 + # compute offsets once per body (avoids 12 redundant tree-ancestry walks in jac_dof) + offset1 = pos1 - wp.vec3(subtree_com_in[worldid, body_rootid[b1]]) + offset2 = pos2 - wp.vec3(subtree_com_in[worldid, body_rootid[b2]]) - # TODO(team): jacdif + # body1 DOFs: b1 is in subtree, b2 is not -> jacdif = 0 - jacp1 = -jacp1 + for k in range(3): + cdof = cdof_in[worldid, dofi + k] + cdof_ang = wp.spatial_top(cdof) + cdof_lin = wp.spatial_bottom(cdof) + jacp1 = cdof_lin + wp.cross(cdof_ang, offset1) + flexedge_J_out[worldid, rowadr + k] = wp.dot(-jacp1, edge) - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi0, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi0, worldid) - jacdif = jacp2 - jacp1 - Ji0 = wp.dot(jacdif, edge) - - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi1, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi1, worldid) - jacdif = jacp2 - jacp1 - Ji1 = wp.dot(jacdif, edge) - - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi2, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi2, worldid) - jacdif = jacp2 - jacp1 - Ji2 = wp.dot(jacdif, edge) - - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj0, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj0, worldid) - jacdif = jacp2 - jacp1 - Jj0 = wp.dot(jacdif, edge) - - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj1, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj1, worldid) - jacdif = jacp2 - jacp1 - Jj1 = wp.dot(jacdif, edge) - - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj2, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj2, worldid) - jacdif = jacp2 - jacp1 - Jj2 = wp.dot(jacdif, edge) - - flexedge_J_out[worldid, 0, sparseid0] = Ji0 - flexedge_J_out[worldid, 0, sparseid1] = Ji1 - flexedge_J_out[worldid, 0, sparseid2] = Ji2 - flexedge_J_out[worldid, 0, sparseid3] = Jj0 - flexedge_J_out[worldid, 0, sparseid4] = Jj1 - flexedge_J_out[worldid, 0, sparseid5] = Jj2 + # body2 DOFs: b2 is in subtree, b1 is not -> jacdif = jacp2 - 0 = jacp2 + for k in range(3): + cdof = cdof_in[worldid, dofj + k] + cdof_ang = wp.spatial_top(cdof) + cdof_lin = wp.spatial_bottom(cdof) + jacp2 = cdof_lin + wp.cross(cdof_ang, offset2) + flexedge_J_out[worldid, rowadr + 3 + k] = wp.dot(jacp2, edge) @event_scope @@ -414,10 +388,8 @@ def flex(m: Model, d: Data): dim=(d.nworld, m.nflexedge), inputs=[ m.nflex, - m.body_parentid, m.body_rootid, m.body_dofadr, - m.dof_bodyid, m.flex_vertadr, m.flex_edgeadr, m.flex_edgenum, @@ -818,7 +790,9 @@ def _qM_sparse( bodyid = dof_bodyid[dofid] # init M(i,i) with armature inertia - qM_out[worldid, 0, madr_ij] = dof_armature[worldid, dofid] + qM_out[worldid, 0, madr_ij] = dof_armature[ + worldid % dof_armature.shape[0], dofid + ] # precompute buf = crb_body_i * cdof_i buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid]) @@ -907,7 +881,7 @@ def _tendon_armature( if is_sparse: # is_sparse is not batched madr_ij = dof_Madr[dofid] - armature = tendon_armature[worldid, tenid] + armature = tendon_armature[worldid % tendon_armature.shape[0], tenid] if armature == 0.0: return @@ -1320,7 +1294,7 @@ def _cfrc_ext_equality( ) id = efc_id_in[worldid, efcid] - eq_data_ = eq_data[worldid, id] + eq_data_ = eq_data[worldid % eq_data.shape[0], id] body_semantic = eq_objtype[id] == ObjType.BODY obj1 = eq_obj1id[id] @@ -1341,7 +1315,7 @@ def _cfrc_ext_equality( else: offset = wp.vec3(eq_data_[3], eq_data_[4], eq_data_[5]) else: - offset = site_pos[worldid, obj1] + offset = site_pos[worldid % site_pos.shape[0], obj1] # transform point on body1: local -> global pos = xmat_in[worldid, bodyid1] @ offset + xpos_in[worldid, bodyid1] @@ -1363,7 +1337,7 @@ def _cfrc_ext_equality( else: offset = wp.vec3(eq_data_[0], eq_data_[1], eq_data_[2]) else: - offset = site_pos[worldid, obj2] + offset = site_pos[worldid % site_pos.shape[0], obj2] # transform point on body2: local -> global pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2] @@ -1558,7 +1532,7 @@ def _tendon_dot( ): worldid, tenid = wp.tid() - armature = tendon_armature[worldid, tenid] + armature = tendon_armature[worldid % tendon_armature.shape[0], tenid] if armature == 0.0: return @@ -1716,7 +1690,7 @@ def _tendon_bias_coef( ): worldid, tenid, dofid = wp.tid() - armature = tendon_armature[worldid, tenid] + armature = tendon_armature[worldid % tendon_armature.shape[0], tenid] if armature == 0.0: return @@ -1740,7 +1714,7 @@ def _tendon_bias_qfrc( ): worldid, tenid, dofid = wp.tid() - armature = tendon_armature[worldid, tenid] + armature = tendon_armature[worldid % tendon_armature.shape[0], tenid] if armature == 0.0: return @@ -1914,40 +1888,45 @@ def com_vel(m: Model, d: Data): @wp.kernel 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), - tendon_adr: wp.array(dtype=int), - tendon_num: wp.array(dtype=int), - wrap_type: wp.array(dtype=int), - wrap_objid: 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.array(dtype=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.array3d(dtype=float), - ten_length_in: wp.array2d(dtype=float), - # Data out: - actuator_length_out: wp.array2d(dtype=float), - actuator_moment_out: wp.array3d(dtype=float), + # 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), + tendon_adr: wp.array(dtype=int), + tendon_num: wp.array(dtype=int), + wrap_type: wp.array(dtype=int), + wrap_objid: 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), + # 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.array3d(dtype=float), + ten_length_in: wp.array2d(dtype=float), + # In: + moment_nnz: wp.array(dtype=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), ): worldid, actid = wp.tid() trntype = actuator_trntype[actid] @@ -1960,20 +1939,33 @@ def _transmission( qadr = jnt_qposadr[jntid] vadr = jnt_dofadr[jntid] if jnt_typ == JointType.FREE: + moment_rownnz_out[worldid, actid] = 6 + rowadr = wp.atomic_add(moment_nnz, worldid, 6) + moment_rowadr_out[worldid, actid] = rowadr + moment_colind_out[worldid, rowadr + 0] = vadr + 0 + moment_colind_out[worldid, rowadr + 1] = vadr + 1 + moment_colind_out[worldid, rowadr + 2] = vadr + 2 + moment_colind_out[worldid, rowadr + 3] = vadr + 3 + moment_colind_out[worldid, rowadr + 4] = vadr + 4 + moment_colind_out[worldid, rowadr + 5] = vadr + 5 actuator_length_out[worldid, actid] = 0.0 if trntype == TrnType.JOINTINPARENT: quat = wp.normalize(wp.quat(qpos[qadr + 3], qpos[qadr + 4], qpos[qadr + 5], qpos[qadr + 6])) quat_neg = math.quat_inv(quat) gearaxis = math.rot_vec_quat(wp.spatial_bottom(gear), quat_neg) - actuator_moment_out[worldid, actid, vadr + 0] = gear[0] - actuator_moment_out[worldid, actid, vadr + 1] = gear[1] - actuator_moment_out[worldid, actid, vadr + 2] = gear[2] - actuator_moment_out[worldid, actid, vadr + 3] = gearaxis[0] - actuator_moment_out[worldid, actid, vadr + 4] = gearaxis[1] - actuator_moment_out[worldid, actid, vadr + 5] = gearaxis[2] + actuator_moment_out[worldid, rowadr + 0] = gear[0] + actuator_moment_out[worldid, rowadr + 1] = gear[1] + actuator_moment_out[worldid, rowadr + 2] = gear[2] + actuator_moment_out[worldid, rowadr + 3] = gearaxis[0] + actuator_moment_out[worldid, rowadr + 4] = gearaxis[1] + actuator_moment_out[worldid, rowadr + 5] = gearaxis[2] else: - for i in range(6): - actuator_moment_out[worldid, actid, vadr + i] = gear[i] + actuator_moment_out[worldid, rowadr + 0] = gear[0] + actuator_moment_out[worldid, rowadr + 1] = gear[1] + actuator_moment_out[worldid, rowadr + 2] = gear[2] + actuator_moment_out[worldid, rowadr + 3] = gear[3] + actuator_moment_out[worldid, rowadr + 4] = gear[4] + actuator_moment_out[worldid, rowadr + 5] = gear[5] elif jnt_typ == JointType.BALL: q = wp.quat(qpos[qadr + 0], qpos[qadr + 1], qpos[qadr + 2], qpos[qadr + 3]) q = wp.normalize(q) @@ -1983,11 +1975,25 @@ def _transmission( quat_neg = math.quat_inv(q) gearaxis = math.rot_vec_quat(gearaxis, quat_neg) actuator_length_out[worldid, actid] = wp.dot(axis_angle, gearaxis) + + nnz = 3 + moment_rownnz_out[worldid, actid] = nnz + rowadr = wp.atomic_add(moment_nnz, worldid, nnz) + moment_rowadr_out[worldid, actid] = rowadr + for i in range(3): - actuator_moment_out[worldid, actid, vadr + i] = gearaxis[i] + sparseid = rowadr + i + moment_colind_out[worldid, sparseid] = vadr + i + actuator_moment_out[worldid, sparseid] = gearaxis[i] elif jnt_typ == JointType.SLIDE or jnt_typ == JointType.HINGE: actuator_length_out[worldid, actid] = qpos[qadr] * gear[0] - actuator_moment_out[worldid, actid, vadr] = gear[0] + + nnz = 1 + moment_rownnz_out[worldid, actid] = nnz + rowadr = wp.atomic_add(moment_nnz, worldid, nnz) + moment_rowadr_out[worldid, actid] = rowadr + moment_colind_out[worldid, rowadr] = vadr + actuator_moment_out[worldid, rowadr] = gear[0] else: wp.printf("unrecognized joint type") elif trntype == TrnType.SLIDERCRANK: @@ -1996,7 +2002,7 @@ def _transmission( id = trnid[0] idslider = trnid[1] gear0 = gear[0] - rod = actuator_cranklength[actid] + rod = actuator_cranklength[worldid % actuator_cranklength.shape[0], actid] site_xmat = site_xmat_in[worldid, idslider] axis = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2]) site_xpos_id = site_xpos_in[worldid, id] @@ -2027,26 +2033,77 @@ def _transmission( dldv = axis dlda = vec - # apply chain rule - # TODO(team): parallelize? - for i in range(nv): + # count dofs + b1 = body_weldid[site_bodyid[id]] + b2 = body_weldid[site_bodyid[idslider]] + da1_init = int(-1) + da2_init = int(-1) + if b1 > 0: + da1_init = body_dofadr[b1] + body_dofnum[b1] - 1 + if b2 > 0: + da2_init = body_dofadr[b2] + body_dofnum[b2] - 1 + + da1 = da1_init + da2 = da2_init + ndof = int(0) + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + ndof += 1 + if da1 == da: + da1 = dof_parentid[da1] + if da2 == da: + da2 = dof_parentid[da2] + + moment_rownnz_out[worldid, actid] = ndof + rowadr = wp.atomic_add(moment_nnz, worldid, ndof) + moment_rowadr_out[worldid, actid] = rowadr + + # traverse dofs + da1 = da1_init + da2 = da2_init + + ptr = ndof - 1 + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + # get Jacobians of axis(jacA) and vec(jac) - # mj_jacPointAxis jacp, jacr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, site_xpos_idslider, site_bodyid[idslider], i, worldid + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + site_xpos_idslider, + site_bodyid[idslider], + da, + worldid, ) jacS = jacp jacA = wp.cross(jacr, axis) - - # mj_jacSite jac, _ = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, site_xpos_id, site_bodyid[id], i, worldid + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + site_xpos_id, + site_bodyid[id], + da, + worldid, ) jac -= jacS # apply the chain rule moment = wp.dot(dlda, jacA) + wp.dot(dldv, jac) - actuator_moment_out[worldid, actid, i] = moment * gear0 + sparseid = rowadr + ptr + moment_colind_out[worldid, sparseid] = da + actuator_moment_out[worldid, sparseid] = moment * gear0 + ptr -= 1 + + if da1 == da: + da1 = dof_parentid[da1] + if da2 == da: + da2 = dof_parentid[da2] elif trntype == TrnType.TENDON: tenid = actuator_trnid[actid][0] @@ -2057,19 +2114,46 @@ def _transmission( adr = tendon_adr[tenid] if wrap_type[adr] == WrapType.JOINT: ten_num = tendon_num[tenid] + rowadr = wp.atomic_add(moment_nnz, worldid, ten_num) + moment_rownnz_out[worldid, actid] = ten_num + moment_rowadr_out[worldid, actid] = rowadr + for i in range(ten_num): dofadr = jnt_dofadr[wrap_objid[adr + i]] - actuator_moment_out[worldid, actid, dofadr] = ten_J_in[worldid, tenid, dofadr] * gear0 + sparseid = rowadr + i + moment_colind_out[worldid, sparseid] = dofadr + actuator_moment_out[worldid, sparseid] = ( + ten_J_in[worldid, tenid, dofadr] * gear0 + ) else: # spatial + # TODO(team): sparse tendon jacobian + ten_nnz = int(0) for dofadr in range(nv): - actuator_moment_out[worldid, actid, dofadr] = ten_J_in[worldid, tenid, dofadr] * gear0 + if ten_J_in[worldid, tenid, dofadr] != 0.0: + ten_nnz += 1 + rowadr = wp.atomic_add(moment_nnz, worldid, ten_nnz) + moment_rownnz_out[worldid, actid] = ten_nnz + moment_rowadr_out[worldid, actid] = rowadr + ptr = int(0) + for dofadr in range(nv): + J = ten_J_in[worldid, tenid, dofadr] + if J != 0.0: + sparseid = rowadr + ptr + moment_colind_out[worldid, sparseid] = dofadr + actuator_moment_out[worldid, sparseid] = J * gear0 + ptr += 1 elif trntype == TrnType.BODY: # cannot compute meaningful length, set to zero actuator_length_out[worldid, actid] = 0.0 # initialize moment + rowadr = wp.atomic_add(moment_nnz, worldid, nv) + moment_rownnz_out[worldid, actid] = nv + moment_rowadr_out[worldid, actid] = rowadr for i in range(nv): - actuator_moment_out[worldid, actid, i] = 0.0 + sparseid = rowadr + i + moment_colind_out[worldid, sparseid] = i + actuator_moment_out[worldid, sparseid] = 0.0 # moment computed by _transmission_body_moment and _transmission_body_moment_scale elif trntype == TrnType.SITE: @@ -2077,7 +2161,7 @@ def _transmission( siteid = trnid[0] refid = trnid[1] - gear = actuator_gear[worldid, actid] + gear = actuator_gear[actuator_gear_id, actid] site_quat_id = worldid % site_quat.shape[0] gear_translation = wp.spatial_top(gear) gear_rotational = wp.spatial_bottom(gear) @@ -2089,22 +2173,46 @@ def _transmission( wrench_translation = site_xmat @ gear_translation wrench_rotation = site_xmat @ gear_rotational - # moment: global Jacobian projected on wrench - # TODO(team): parallelize - for i in range(nv): + # count dofs + b1 = body_weldid[site_bodyid[siteid]] + da_init = int(-1) + if b1 > 0: + da_init = body_dofadr[b1] + body_dofnum[b1] - 1 + + da = da_init + ndof = int(0) + while da >= 0: + ndof += 1 + da = dof_parentid[da] + + moment_rownnz_out[worldid, actid] = ndof + rowadr = wp.atomic_add(moment_nnz, worldid, ndof) + moment_rowadr_out[worldid, actid] = rowadr + actuator_length_out[worldid, actid] = 0.0 + + # traverse dofs + da = da_init + ptr = ndof - 1 + while da >= 0: jacp, jacr = support.jac_dof( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - site_xpos_in[worldid, siteid], - site_bodyid[siteid], - i, - worldid, + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + site_xpos_in[worldid, siteid], + site_bodyid[siteid], + da, + worldid, ) - actuator_length_out[worldid, actid] = 0.0 - actuator_moment_out[worldid, actid, i] = wp.dot(jacp, wrench_translation) + wp.dot(jacr, wrench_rotation) + moment = wp.dot(jacp, wrench_translation) + wp.dot( + jacr, wrench_rotation + ) + sparseid = rowadr + ptr + moment_colind_out[worldid, sparseid] = da + actuator_moment_out[worldid, sparseid] = moment + ptr -= 1 + da = dof_parentid[da] # reference site defined else: # initialize last dof address for each body @@ -2162,69 +2270,116 @@ def _transmission( actuator_length_out[worldid, actid] = length - # TODO(team): parallelize - for i in range(nv): + # count dofs + da1_init = int(-1) + da2_init = int(-1) + if b0 > 0: + da1_init = body_dofadr[b0] + body_dofnum[b0] - 1 + if b1 > 0: + da2_init = body_dofadr[b1] + body_dofnum[b1] - 1 + + da1 = da1_init + da2 = da2_init + ndof = int(0) + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + if da1 == da and da2 == da: + break + ndof += 1 + if da1 == da: + da1 = dof_parentid[da1] + if da2 == da: + da2 = dof_parentid[da2] + + moment_rownnz_out[worldid, actid] = ndof + rowadr = wp.atomic_add(moment_nnz, worldid, ndof) + moment_rowadr_out[worldid, actid] = rowadr + + # traverse dofs + da1 = da1_init + da2 = da2_init + + ptr = ndof - 1 + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + if da1 == da and da2 == da: + break + jacp, jacr = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, site_xpos, site_bodyid[siteid], i, worldid + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + site_xpos, + site_bodyid[siteid], + da, + worldid, ) - - # jacref: global Jacobian of reference site jacpref, jacrref = support.jac_dof( - body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, ref_xpos, site_bodyid[refid], i, worldid + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + ref_xpos, + site_bodyid[refid], + da, + worldid, ) - jacpdif = jacp - jacpref - jacrdif = jacr - jacrref - - # if common ancestral dof was found, clear the columns of its parental chain - da = dofadr_common - while da >= 0: - if da == i: - jacpdif = wp.vec3(0.0) - jacrdif = wp.vec3(0.0) - break - da = dof_parentid[da] - - # moment: global Jacobian projected on wrench moment = float(0.0) - if translational_transmission: - moment += wp.dot(jacpdif, wrench_translation) + moment += wp.dot(jacp - jacpref, wrench_translation) if rotational_transmission: - moment += wp.dot(jacrdif, wrench_rotation) + moment += wp.dot(jacr - jacrref, wrench_rotation) - actuator_moment_out[worldid, actid, i] = moment + sparseid = rowadr + ptr + moment_colind_out[worldid, sparseid] = da + actuator_moment_out[worldid, sparseid] = moment + ptr -= 1 + + if da1 == da: + da1 = dof_parentid[da1] + if da2 == da: + da2 = dof_parentid[da2] else: wp.printf("unhandled transmission type %d\n", trntype) @wp.kernel 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), - # Data in: - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_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_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_in: wp.array3d(dtype=float), - nacon_in: wp.array(dtype=int), - # Data out: - actuator_moment_out: wp.array3d(dtype=float), - # Out: - actuator_trntype_body_ncon_out: wp.array2d(dtype=int), + # 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), + # 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), + # In: + efc_is_sparse: bool, + # Data out: + actuator_moment_out: wp.array2d(dtype=float), + # Out: + actuator_trntype_body_ncon_out: wp.array2d(dtype=int), ): trnbodyid, conid, dofid = wp.tid() actid = actuator_trntype_body_adr[trnbodyid] @@ -2257,23 +2412,61 @@ def _transmission_body_moment( if dofid == 0: wp.atomic_add(actuator_trntype_body_ncon_out[worldid], trnbodyid, 1) + rowadr = moment_rowadr_in[worldid, actid] + # mark contact normals in efc_force if contact_exclude == 0: contact_dim = contact_dim_in[conid] contact_efc_address = contact_efc_address_in[conid] if contact_dim == 1 or opt_cone == ConeType.ELLIPTIC: - efc_force = 1.0 efcid0 = contact_efc_address[0] - wp.atomic_add(actuator_moment_out[worldid, actid], dofid, efc_J_in[worldid, efcid0, dofid] * efc_force) - + if efc_is_sparse: + rownnz = efc_J_rownnz_in[worldid, efcid0] + if dofid < rownnz: + efc_rowadr = efc_J_rowadr_in[worldid, efcid0] + efc_sparseid = efc_rowadr + dofid + colind = efc_J_colind_in[worldid, 0, efc_sparseid] + wp.atomic_add( + actuator_moment_out[worldid], + rowadr + colind, + efc_J_in[worldid, 0, efc_sparseid], + ) + else: + return + else: + colind = dofid + wp.atomic_add( + actuator_moment_out[worldid], + rowadr + colind, + efc_J_in[worldid, efcid0, dofid], + ) else: npyramid = contact_dim - 1 # number of frictional directions efc_force = 0.5 / float(npyramid) for j in range(2 * npyramid): efcid = contact_efc_address[j] - wp.atomic_add(actuator_moment_out[worldid, actid], dofid, efc_J_in[worldid, efcid, dofid] * efc_force) + if efc_is_sparse: + rownnz = efc_J_rownnz_in[worldid, efcid] + if dofid < rownnz: + efc_rowadr = efc_J_rowadr_in[worldid, efcid] + efc_sparseid = efc_rowadr + dofid + colind = efc_J_colind_in[worldid, 0, efc_sparseid] + wp.atomic_add( + actuator_moment_out[worldid], + rowadr + colind, + efc_J_in[worldid, 0, efc_sparseid] * efc_force, + ) + else: + return + else: + colind = dofid + wp.atomic_add( + actuator_moment_out[worldid], + rowadr + colind, + efc_J_in[worldid, efcid, dofid] * efc_force, + ) # excluded contact in gap: get Jacobian, accumulate elif contact_exclude == 1: @@ -2282,22 +2475,58 @@ def _transmission_body_moment( normal = wp.vec3(contact_frame[0, 0], contact_frame[0, 1], contact_frame[0, 2]) # get Jacobian difference - jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, contact_pos, b1, dofid, worldid) - jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, contact_pos, b2, dofid, worldid) + efcid0 = contact_efc_address_in[conid][0] + if efc_is_sparse and efcid0 >= 0: + # contact has valid efc row: use sparse pattern + if dofid >= efc_J_rownnz_in[worldid, efcid0]: + return + sparseid = efc_J_rowadr_in[worldid, efcid0] + dofid + colind = efc_J_colind_in[worldid, 0, sparseid] + else: + # excluded contact with no efc row or dense: use dofid directly + colind = dofid + + jacp1, _ = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + contact_pos, + b1, + colind, + worldid, + ) + jacp2, _ = support.jac_dof( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + contact_pos, + b2, + colind, + worldid, + ) + jacdif = jacp2 - jacp1 # project Jacobian along the normal of the contact frame - wp.atomic_add(actuator_moment_out[worldid, actid], dofid, wp.dot(normal, jacdif)) + wp.atomic_add( + actuator_moment_out[worldid], rowadr + colind, wp.dot(normal, jacdif) + ) @wp.kernel def _transmission_body_moment_scale( - # Model: - actuator_trntype_body_adr: wp.array(dtype=int), - # In: - actuator_trntype_body_ncon_in: wp.array2d(dtype=int), - # Data out: - actuator_moment_out: wp.array3d(dtype=float), + # Model: + actuator_trntype_body_adr: wp.array(dtype=int), + # Data in: + moment_rowadr_in: wp.array2d(dtype=int), + # In: + actuator_trntype_body_ncon_in: wp.array2d(dtype=int), + # Data out: + actuator_moment_out: wp.array2d(dtype=float), ): worldid, trnbodyid, dofid = wp.tid() @@ -2305,7 +2534,8 @@ def _transmission_body_moment_scale( if ncon > 0: actid = actuator_trntype_body_adr[trnbodyid] - actuator_moment_out[worldid, actid, dofid] /= -float(ncon) + rowadr = moment_rowadr_in[worldid, actid] + actuator_moment_out[worldid, rowadr + dofid] /= -float(ncon) @event_scope @@ -2315,80 +2545,95 @@ def transmission(m: Model, d: Data): Updates the actuator length and moments for all actuators in the model, including joint and tendon transmissions. """ - d.actuator_moment.zero_() + # TODO(team): investigate pre-computing moment_rownnz, moment_rowadr, moment_colind + moment_nnz = wp.zeros((d.nworld,), dtype=int) + wp.launch( - _transmission, - dim=[d.nworld, m.nu], - inputs=[ - m.nv, - m.body_parentid, - m.body_rootid, - m.body_weldid, - m.body_dofnum, - m.body_dofadr, - m.jnt_type, - m.jnt_qposadr, - m.jnt_dofadr, - m.dof_bodyid, - m.dof_parentid, - m.site_bodyid, - m.site_quat, - m.tendon_adr, - m.tendon_num, - m.wrap_type, - m.wrap_objid, - m.actuator_trntype, - m.actuator_trnid, - m.actuator_gear, - m.actuator_cranklength, - d.qpos, - d.xquat, - d.site_xpos, - d.site_xmat, - d.subtree_com, - d.cdof, - d.ten_J, - d.ten_length, - ], - outputs=[d.actuator_length, d.actuator_moment], + _transmission, + dim=(d.nworld, m.nu), + inputs=[ + m.nv, + m.body_parentid, + m.body_rootid, + m.body_weldid, + m.body_dofnum, + m.body_dofadr, + m.jnt_type, + m.jnt_qposadr, + m.jnt_dofadr, + m.dof_bodyid, + m.dof_parentid, + m.site_bodyid, + m.site_quat, + m.tendon_adr, + m.tendon_num, + m.wrap_type, + m.wrap_objid, + m.actuator_trntype, + m.actuator_trnid, + m.actuator_gear, + m.actuator_cranklength, + d.qpos, + d.xquat, + d.site_xpos, + d.site_xmat, + d.subtree_com, + d.cdof, + d.ten_J, + d.ten_length, + moment_nnz, + ], + outputs=[ + d.actuator_length, + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + d.actuator_moment, + ], ) if m.nacttrnbody: # compute moments ncon = wp.zeros((d.nworld, m.nacttrnbody), dtype=int) + wp.launch( - _transmission_body_moment, - dim=(m.nacttrnbody, d.naconmax, m.nv), - inputs=[ - m.opt.cone, - m.body_parentid, - m.body_rootid, - m.dof_bodyid, - m.geom_bodyid, - m.actuator_trnid, - m.actuator_trntype_body_adr, - d.subtree_com, - d.cdof, - d.contact.dist, - d.contact.pos, - d.contact.frame, - d.contact.includemargin, - d.contact.dim, - d.contact.geom, - d.contact.efc_address, - d.contact.worldid, - d.efc.J, - d.nacon, - ], - outputs=[d.actuator_moment, ncon], + _transmission_body_moment, + dim=(m.nacttrnbody, d.naconmax, m.nv), + inputs=[ + m.opt.cone, + m.body_parentid, + m.body_rootid, + m.dof_bodyid, + m.geom_bodyid, + m.actuator_trnid, + m.actuator_trntype_body_adr, + d.subtree_com, + d.cdof, + d.moment_rowadr, + d.contact.dist, + d.contact.pos, + d.contact.frame, + d.contact.includemargin, + d.contact.dim, + d.contact.geom, + d.contact.efc_address, + d.contact.worldid, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.nacon, + SPARSE_CONSTRAINT_JACOBIAN, + ], + outputs=[d.actuator_moment, ncon], ) # scale moments wp.launch( - _transmission_body_moment_scale, - dim=(d.nworld, m.nacttrnbody, m.nv), - inputs=[m.actuator_trntype_body_adr, ncon], - outputs=[d.actuator_moment], + _transmission_body_moment_scale, + dim=(d.nworld, m.nacttrnbody, m.nv), + inputs=[m.actuator_trntype_body_adr, d.moment_rowadr, ncon], + outputs=[d.actuator_moment], ) @@ -2888,7 +3133,7 @@ def _spatial_geom_tendon( # get geom information geom_xpos = geom_xpos_in[worldid, wrap_objid_geom] geom_xmat = geom_xmat_in[worldid, wrap_objid_geom] - geomsize = geom_size[worldid, wrap_objid_geom][0] + geomsize = geom_size[worldid % geom_size.shape[0], wrap_objid_geom][0] geom_type = wrap_type[wrap_adr] # get body ids for site-geom-site instances @@ -3147,6 +3392,7 @@ def _spatial_tendon_wrap( ten_wrapnum_out[worldid, i] = wrapnum +@event_scope def tendon(m: Model, d: Data): """Computes tendon lengths and moments. 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 8eabb52f..82ac23c7 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -17,16 +17,17 @@ import dataclasses from math import ceil from math import sqrt -import warp as wp - from mujoco.mjx.third_party.mujoco_warp._src import math from mujoco.mjx.third_party.mujoco_warp._src import smooth from mujoco.mjx.third_party.mujoco_warp._src import support from mujoco.mjx.third_party.mujoco_warp._src import types from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_func from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func +from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope +from mujoco.mjx.third_party.mujoco_warp._src.warp_util import scoped_mathdx_gemm_disabled +import warp as wp wp.set_module_options({"enable_backward": False}) @@ -43,6 +44,8 @@ class InverseContext: 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) @dataclasses.dataclass @@ -69,6 +72,9 @@ class SolverContext: beta: wp.array(dtype=float) h: wp.array3d(dtype=float) hfactor: wp.array3d(dtype=float) + # Incremental Hessian update (Newton only) + changed_efc_ids: wp.array2d(dtype=int) + changed_efc_count: wp.array(dtype=int) def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext: @@ -85,12 +91,14 @@ def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext: njmax = d.njmax return InverseContext( - Jaref=wp.empty((nworld, njmax), dtype=float), - search_dot=wp.empty((nworld,), dtype=float), - gauss=wp.empty((nworld,), dtype=float), - cost=wp.empty((nworld,), dtype=float), - prev_cost=wp.empty((nworld,), dtype=float), - done=wp.empty((nworld,), dtype=bool), + Jaref=wp.empty((nworld, njmax), dtype=float), + search_dot=wp.empty((nworld,), dtype=float), + gauss=wp.empty((nworld,), dtype=float), + cost=wp.empty((nworld,), dtype=float), + prev_cost=wp.empty((nworld,), dtype=float), + done=wp.empty((nworld,), dtype=bool), + changed_efc_ids=wp.empty((nworld, 0), dtype=int), + changed_efc_count=wp.empty((0,), dtype=int), ) @@ -109,31 +117,40 @@ def create_solver_context(m: types.Model, d: types.Data) -> SolverContext: nv_pad = m.nv_pad njmax = d.njmax - # Newton solver needs h; hfactor only needed if nv > _BLOCK_CHOLESKY_DIM alloc_h = m.opt.solver == types.SolverType.NEWTON alloc_hfactor = alloc_h and nv > _BLOCK_CHOLESKY_DIM return SolverContext( - Jaref=wp.empty((nworld, njmax), dtype=float), - search_dot=wp.empty((nworld,), dtype=float), - gauss=wp.empty((nworld,), dtype=float), - cost=wp.empty((nworld,), dtype=float), - prev_cost=wp.empty((nworld,), dtype=float), - done=wp.empty((nworld,), dtype=bool), - grad=wp.zeros((nworld, nv_pad), dtype=float), - grad_dot=wp.empty((nworld,), dtype=float), - Mgrad=wp.zeros((nworld, nv_pad), dtype=float), - search=wp.empty((nworld, nv), dtype=float), - mv=wp.empty((nworld, nv), dtype=float), - jv=wp.empty((nworld, njmax), dtype=float), - quad=wp.empty((nworld, njmax), dtype=wp.vec3), - quad_gauss=wp.empty((nworld,), dtype=wp.vec3), - alpha=wp.empty((nworld,), dtype=float), - prev_grad=wp.empty((nworld, nv), dtype=float), - prev_Mgrad=wp.empty((nworld, nv), dtype=float), - beta=wp.empty((nworld,), dtype=float), - h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float), - hfactor=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_hfactor else wp.empty((nworld, 0, 0), dtype=float), + Jaref=wp.empty((nworld, njmax), dtype=float), + search_dot=wp.empty((nworld,), dtype=float), + gauss=wp.empty((nworld,), dtype=float), + cost=wp.empty((nworld,), dtype=float), + prev_cost=wp.empty((nworld,), dtype=float), + done=wp.empty((nworld,), dtype=bool), + grad=wp.zeros((nworld, nv_pad), dtype=float), + grad_dot=wp.empty((nworld,), dtype=float), + Mgrad=wp.zeros((nworld, nv_pad), dtype=float), + search=wp.empty((nworld, nv), dtype=float), + mv=wp.empty((nworld, nv), dtype=float), + jv=wp.empty((nworld, njmax), dtype=float), + quad=wp.empty((nworld, njmax), dtype=wp.vec3), + quad_gauss=wp.empty((nworld,), dtype=wp.vec3), + alpha=wp.empty((nworld,), dtype=float), + prev_grad=wp.empty((nworld, nv), dtype=float), + prev_Mgrad=wp.empty((nworld, nv), dtype=float), + beta=wp.empty((nworld,), dtype=float), + h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) + if alloc_h + else wp.empty((nworld, 0, 0), dtype=float), + hfactor=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) + if alloc_hfactor + else wp.empty((nworld, 0, 0), dtype=float), + changed_efc_ids=wp.empty((nworld, njmax), dtype=int) + if alloc_h + else wp.empty((nworld, 0), dtype=int), + changed_efc_count=wp.empty((nworld,), dtype=int) + if alloc_h + else wp.empty((0,), dtype=int), ) @@ -905,43 +922,46 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: @wp.kernel(module="unique", enable_backward=False) 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), - # 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_in: wp.array3d(dtype=float), - efc_D_in: wp.array2d(dtype=float), - efc_frictionloss_in: wp.array2d(dtype=float), - njmax_in: int, - nacon_in: wp.array(dtype=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), - # Data out: - qacc_out: wp.array2d(dtype=float), - efc_Ma_out: wp.array2d(dtype=float), - # Out: - ctx_Jaref_out: wp.array2d(dtype=float), - ctx_jv_out: wp.array2d(dtype=float), - ctx_quad_out: wp.array2d(dtype=wp.vec3), + # 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), + # 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), + njmax_in: int, + nacon_in: wp.array(dtype=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), + # Data out: + qacc_out: wp.array2d(dtype=float), + efc_Ma_out: wp.array2d(dtype=float), + # Out: + ctx_Jaref_out: wp.array2d(dtype=float), + ctx_jv_out: wp.array2d(dtype=float), + ctx_quad_out: wp.array2d(dtype=wp.vec3), ): worldid, tid = wp.tid() @@ -956,8 +976,18 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: if wp.static(FUSE_JV): for efcid in range(tid, nefc, wp.block_dim()): jv = float(0.0) - for i in range(nv): - jv += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i] + if wp.static(SPARSE_CONSTRAINT_JACOBIAN): + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + for k in range(rownnz): + sparseid = rowadr + k + colind = efc_J_colind_in[worldid, 0, sparseid] + jv += ( + efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind] + ) + else: + for i in range(nv): + jv += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i] ctx_jv_out[worldid, efcid] = jv _syncthreads() # ensure all jv values are written before reading @@ -1031,7 +1061,7 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: snorm = wp.sqrt(ctx_search_dot_in[worldid]) meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]] scale = meaninertia * wp.float(nv) - gtol = tolerance * ls_tolerance * snorm * scale + gtol = wp.max(tolerance * ls_tolerance * snorm * scale, 1e-6) # p0 via parallel reduction local_p0 = wp.vec3(0.0) @@ -1164,8 +1194,8 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: lo_in_sum = wp.tile_reduce(wp.add, lo_in_tile) lo_in = _eval_pt(ctx_quad_gauss, lo_alpha_in) + lo_in_sum[0] - # check for initial convergence: if |derivative| < gtol, accept Newton step immediately - initial_converged = wp.abs(lo_in[1]) < gtol + # accept Newton step if derivative is small and cost improved + initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < p0[0] # main iterative loop - skip if already converged if not initial_converged: @@ -1330,39 +1360,42 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus fuse_jv: Whether jv is computed in-kernel (True) or pre-computed (False). """ wp.launch_tiled( - linesearch_iterative(m.opt.ls_iterations, m.opt.cone, fuse_jv), - dim=d.nworld, - inputs=[ - m.nv, - m.opt.tolerance, - m.opt.ls_tolerance, - m.opt.impratio_invsqrt, - m.stat.meaninertia, - d.ne, - d.nf, - d.nefc, - d.qfrc_smooth, - d.contact.friction, - d.contact.dim, - d.contact.efc_address, - d.efc.type, - d.efc.id, - d.efc.J, - d.efc.D, - d.efc.frictionloss, - d.njmax, - d.nacon, - ctx.Jaref, - ctx.search, - ctx.search_dot, - ctx.gauss, - ctx.mv, - ctx.jv, - ctx.quad, - ctx.done, - ], - outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad], - block_dim=m.block_dim.linesearch_iterative, + linesearch_iterative(m.opt.ls_iterations, m.opt.cone, fuse_jv), + dim=d.nworld, + inputs=[ + m.nv, + m.opt.tolerance, + m.opt.ls_tolerance, + m.opt.impratio_invsqrt, + m.stat.meaninertia, + d.ne, + d.nf, + d.nefc, + d.qfrc_smooth, + d.contact.friction, + d.contact.dim, + d.contact.efc_address, + d.efc.type, + d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.D, + d.efc.frictionloss, + d.njmax, + d.nacon, + ctx.Jaref, + ctx.search, + ctx.search_dot, + ctx.gauss, + ctx.mv, + ctx.jv, + ctx.quad, + ctx.done, + ], + outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad], + block_dim=m.block_dim.linesearch_iterative, ) @@ -1387,17 +1420,20 @@ def linesearch_zero_jv( @cache_kernel -def linesearch_jv_fused(nv: int, dofs_per_thread: int): +def linesearch_jv_fused(opt_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_in: wp.array3d(dtype=float), - # In: - ctx_search_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), - # Out: - ctx_jv_out: wp.array2d(dtype=float), + # 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), + # In: + ctx_search_in: wp.array2d(dtype=float), + ctx_done_in: wp.array(dtype=bool), + # Out: + ctx_jv_out: wp.array2d(dtype=float), ): worldid, efcid, dofstart = wp.tid() @@ -1410,16 +1446,40 @@ def linesearch_jv_fused(nv: int, dofs_per_thread: int): jv_out = float(0.0) if wp.static(dofs_per_thread >= nv): - for i in range(wp.static(min(dofs_per_thread, nv))): - jv_out += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i] + if wp.static(SPARSE_CONSTRAINT_JACOBIAN): + # Sparse: iterate over non-zero entries in the row + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + for k in range(rownnz): + sparseid = rowadr + k + colind = efc_J_colind_in[worldid, 0, sparseid] + jv_out += ( + efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind] + ) + else: + for i in range(wp.static(min(dofs_per_thread, nv))): + jv_out += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i] ctx_jv_out[worldid, efcid] = jv_out else: - for i in range(wp.static(dofs_per_thread)): - ii = dofstart * wp.static(dofs_per_thread) + i - if ii < nv: - jv_out += efc_J_in[worldid, efcid, ii] * ctx_search_in[worldid, ii] - wp.atomic_add(ctx_jv_out, worldid, efcid, jv_out) + if wp.static(SPARSE_CONSTRAINT_JACOBIAN): + # Sparse: thread 0 handles entire row (sparse entries << nv typically) + if dofstart == 0: + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + for k in range(rownnz): + sparseid = rowadr + k + colind = efc_J_colind_in[worldid, 0, sparseid] + jv_out += ( + efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind] + ) + ctx_jv_out[worldid, efcid] = jv_out + else: + for i in range(wp.static(dofs_per_thread)): + ii = dofstart * wp.static(dofs_per_thread) + i + if ii < nv: + jv_out += efc_J_in[worldid, efcid, ii] * ctx_search_in[worldid, ii] + wp.atomic_add(ctx_jv_out, worldid, efcid, jv_out) return kernel @@ -1523,7 +1583,10 @@ def linesearch_prepare_quad( dim = contact_dim_in[conid] friction = contact_friction_in[conid] - mu = friction[0] * opt_impratio_invsqrt[worldid] + mu = ( + friction[0] + * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] + ) u0 = Jaref * mu v0 = jv * mu @@ -1625,9 +1688,10 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra # mv = qM @ search (common to both parallel and iterative) support.mul_m(m, d, ctx.mv, ctx.search, skip=ctx.done) - # Fuse jv computation in-kernel for small nv (iterative only) + # Fuse jv computation in-kernel for small nv (iterative only, dense only) # Parallel linesearch always requires jv pre-computed - fuse_jv = m.nv <= 50 and not m.opt.ls_parallel + # Sparse mode requires pre-computed jv since in-kernel uses dense indexing + fuse_jv = m.nv <= 50 and not m.opt.ls_parallel and not m.is_sparse # jv = J @ search (when not fused into iterative kernel) if not fuse_jv: @@ -1643,10 +1707,18 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra ) wp.launch( - linesearch_jv_fused(m.nv, dofs_per_thread), - dim=(d.nworld, d.njmax, threads_per_efc), - inputs=[d.nefc, d.efc.J, ctx.search, ctx.done], - outputs=[ctx.jv], + linesearch_jv_fused(m.is_sparse, m.nv, dofs_per_thread), + dim=(d.nworld, d.njmax, threads_per_efc), + inputs=[ + d.nefc, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + ctx.search, + ctx.done, + ], + outputs=[ctx.jv], ) if m.opt.ls_parallel: @@ -1672,16 +1744,19 @@ def solve_init_efc( @cache_kernel -def solve_init_jaref(nv: int, dofs_per_thread: int): +def solve_init_jaref(opt_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_in: wp.array3d(dtype=float), - efc_aref_in: wp.array2d(dtype=float), - # Out: - ctx_Jaref_out: wp.array2d(dtype=float), + # 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), + # Out: + ctx_Jaref_out: wp.array2d(dtype=float), ): worldid, efcid, dofstart = wp.tid() @@ -1689,22 +1764,32 @@ def solve_init_jaref(nv: int, dofs_per_thread: int): return jaref = float(0.0) - - if wp.static(dofs_per_thread >= nv): - for i in range(wp.static(min(dofs_per_thread, nv))): - jaref += efc_J_in[worldid, efcid, i] * qacc_in[worldid, i] + if wp.static(SPARSE_CONSTRAINT_JACOBIAN): + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + for i in range(rownnz): + sparseid = rowadr + i + colind = efc_J_colind_in[worldid, 0, sparseid] + jaref += efc_J_in[worldid, 0, sparseid] * qacc_in[worldid, colind] ctx_Jaref_out[worldid, efcid] = jaref - efc_aref_in[worldid, efcid] - else: - for i in range(wp.static(dofs_per_thread)): - ii = dofstart * wp.static(dofs_per_thread) + i - if ii < nv: - jaref += efc_J_in[worldid, efcid, ii] * qacc_in[worldid, ii] + if wp.static(dofs_per_thread >= nv): + for i in range(wp.static(min(dofs_per_thread, nv))): + jaref += efc_J_in[worldid, efcid, i] * qacc_in[worldid, i] + ctx_Jaref_out[worldid, efcid] = jaref - efc_aref_in[worldid, efcid] - if dofstart == 0: - wp.atomic_add(ctx_Jaref_out, worldid, efcid, jaref - efc_aref_in[worldid, efcid]) else: - wp.atomic_add(ctx_Jaref_out, worldid, efcid, jaref) + for i in range(wp.static(dofs_per_thread)): + ii = dofstart * wp.static(dofs_per_thread) + i + if ii < nv: + jaref += efc_J_in[worldid, efcid, ii] * qacc_in[worldid, ii] + + if dofstart == 0: + wp.atomic_add( + ctx_Jaref_out, worldid, efcid, jaref - efc_aref_in[worldid, efcid] + ) + else: + wp.atomic_add(ctx_Jaref_out, worldid, efcid, jaref) return kernel @@ -1743,144 +1828,204 @@ def update_constraint_init_cost( ctx_cost_out[worldid] = 0.0 +@cache_kernel +def update_constraint_efc(track_changes: bool): + TRACK_CHANGES = track_changes + + @wp.kernel(module="unique", enable_backward=False) + def kernel( + # Model: + opt_impratio_invsqrt: wp.array(dtype=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), + # In: + ctx_Jaref_in: wp.array2d(dtype=float), + ctx_done_in: wp.array(dtype=bool), + # Data out: + efc_force_out: wp.array2d(dtype=float), + efc_state_out: wp.array2d(dtype=int), + # Out: + ctx_cost_out: wp.array(dtype=float), + changed_ids_out: wp.array2d(dtype=int), + changed_count_out: wp.array(dtype=int), + ): + worldid, efcid = wp.tid() + + if efcid >= nefc_in[worldid]: + return + + if ctx_done_in[worldid]: + return + + # Read old QUADRATIC status before overwriting + if wp.static(TRACK_CHANGES): + old_quad = ( + efc_state_out[worldid, efcid] == types.ConstraintState.QUADRATIC.value + ) + + efc_D = efc_D_in[worldid, efcid] + Jaref = ctx_Jaref_in[worldid, efcid] + + ne = ne_in[worldid] + nf = nf_in[worldid] + + new_state = types.ConstraintState.SATISFIED.value + + if efcid < ne: + # equality + efc_force_out[worldid, efcid] = -efc_D * Jaref + new_state = types.ConstraintState.QUADRATIC.value + wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) + elif efcid < ne + nf: + # friction + f = efc_frictionloss_in[worldid, efcid] + rf = math.safe_div(f, efc_D) + if Jaref <= -rf: + efc_force_out[worldid, efcid] = f + new_state = types.ConstraintState.LINEARNEG.value + wp.atomic_add(ctx_cost_out, worldid, -f * (0.5 * rf + Jaref)) + elif Jaref >= rf: + efc_force_out[worldid, efcid] = -f + new_state = types.ConstraintState.LINEARPOS.value + wp.atomic_add(ctx_cost_out, worldid, -f * (0.5 * rf - Jaref)) + else: + efc_force_out[worldid, efcid] = -efc_D * Jaref + new_state = types.ConstraintState.QUADRATIC.value + wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) + elif efc_type_in[worldid, efcid] != types.ConstraintType.CONTACT_ELLIPTIC: + # limit, frictionless contact, pyramidal friction cone contact + if Jaref >= 0.0: + efc_force_out[worldid, efcid] = 0.0 + new_state = types.ConstraintState.SATISFIED.value + else: + efc_force_out[worldid, efcid] = -efc_D * Jaref + new_state = types.ConstraintState.QUADRATIC.value + wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) + else: # elliptic friction cone contact + conid = efc_id_in[worldid, efcid] + + if conid >= nacon_in[0]: + return + + dim = contact_dim_in[conid] + friction = contact_friction_in[conid] + mu = ( + friction[0] + * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] + ) + + efcid0 = contact_efc_address_in[conid, 0] + if efcid0 < 0: + return + + N = ctx_Jaref_in[worldid, efcid0] * mu + + ufrictionj = float(0.0) + TT = float(0.0) + for j in range(1, dim): + efcidj = contact_efc_address_in[conid, j] + if efcidj < 0: + return + frictionj = friction[j - 1] + uj = ctx_Jaref_in[worldid, efcidj] * frictionj + TT += uj * uj + if efcid == efcidj: + ufrictionj = uj * frictionj + + if TT <= 0.0: + T = 0.0 + else: + T = wp.sqrt(TT) + + # top zone + if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)): + efc_force_out[worldid, efcid] = 0.0 + new_state = types.ConstraintState.SATISFIED.value + # bottom zone + elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)): + efc_force_out[worldid, efcid] = -efc_D * Jaref + new_state = types.ConstraintState.QUADRATIC.value + wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) + # middle zone + else: + dm = math.safe_div(efc_D_in[worldid, efcid0], mu * mu * (1.0 + mu * mu)) + nmt = N - mu * T + + force = -dm * nmt * mu + + if efcid == efcid0: + efc_force_out[worldid, efcid] = force + wp.atomic_add(ctx_cost_out, worldid, 0.5 * dm * nmt * nmt) + else: + efc_force_out[worldid, efcid] = -math.safe_div(force, T) * ufrictionj + + new_state = types.ConstraintState.CONE.value + + efc_state_out[worldid, efcid] = new_state + + if wp.static(TRACK_CHANGES): + new_quad = new_state == types.ConstraintState.QUADRATIC.value + if old_quad != new_quad: + idx = wp.atomic_add(changed_count_out, worldid, 1) + changed_ids_out[worldid, idx] = efcid + + return kernel + + @wp.kernel -def update_constraint_efc( - # Model: - opt_impratio_invsqrt: wp.array(dtype=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), - # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), - # Data out: - efc_force_out: wp.array2d(dtype=float), - efc_state_out: wp.array2d(dtype=int), - # Out: - ctx_cost_out: wp.array(dtype=float), +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), + # In: + ctx_done_in: wp.array(dtype=bool), + # Data out: + qfrc_constraint_out: wp.array2d(dtype=float), ): worldid, efcid = wp.tid() - if efcid >= nefc_in[worldid]: - return - if ctx_done_in[worldid]: return - efc_D = efc_D_in[worldid, efcid] - Jaref = ctx_Jaref_in[worldid, efcid] + if efcid >= nefc_in[worldid]: + return - ne = ne_in[worldid] - nf = nf_in[worldid] + force = efc_force_in[worldid, efcid] - if efcid < ne: - # equality - efc_force_out[worldid, efcid] = -efc_D * Jaref - efc_state_out[worldid, efcid] = types.ConstraintState.QUADRATIC - wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) - elif efcid < ne + nf: - # friction - f = efc_frictionloss_in[worldid, efcid] - rf = math.safe_div(f, efc_D) - if Jaref <= -rf: - efc_force_out[worldid, efcid] = f - efc_state_out[worldid, efcid] = types.ConstraintState.LINEARNEG - wp.atomic_add(ctx_cost_out, worldid, -f * (0.5 * rf + Jaref)) - elif Jaref >= rf: - efc_force_out[worldid, efcid] = -f - efc_state_out[worldid, efcid] = types.ConstraintState.LINEARPOS - wp.atomic_add(ctx_cost_out, worldid, -f * (0.5 * rf - Jaref)) - else: - efc_force_out[worldid, efcid] = -efc_D * Jaref - efc_state_out[worldid, efcid] = types.ConstraintState.QUADRATIC - wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) - elif efc_type_in[worldid, efcid] != types.ConstraintType.CONTACT_ELLIPTIC: - # limit, frictionless contact, pyramidal friction cone contact - if Jaref >= 0.0: - efc_force_out[worldid, efcid] = 0.0 - efc_state_out[worldid, efcid] = types.ConstraintState.SATISFIED - else: - efc_force_out[worldid, efcid] = -efc_D * Jaref - efc_state_out[worldid, efcid] = types.ConstraintState.QUADRATIC - wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) - else: # elliptic friction cone contact - conid = efc_id_in[worldid, efcid] - - if conid >= nacon_in[0]: - return - - dim = contact_dim_in[conid] - friction = contact_friction_in[conid] - mu = friction[0] * opt_impratio_invsqrt[worldid] - - efcid0 = contact_efc_address_in[conid, 0] - if efcid0 < 0: - return - - N = ctx_Jaref_in[worldid, efcid0] * mu - - ufrictionj = float(0.0) - TT = float(0.0) - for j in range(1, dim): - efcidj = contact_efc_address_in[conid, j] - if efcidj < 0: - return - frictionj = friction[j - 1] - uj = ctx_Jaref_in[worldid, efcidj] * frictionj - TT += uj * uj - if efcid == efcidj: - ufrictionj = uj * frictionj - - if TT <= 0.0: - T = 0.0 - else: - T = wp.sqrt(TT) - - # top zone - if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)): - efc_force_out[worldid, efcid] = 0.0 - efc_state_out[worldid, efcid] = types.ConstraintState.SATISFIED - # bottom zone - elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)): - efc_force_out[worldid, efcid] = -efc_D * Jaref - efc_state_out[worldid, efcid] = types.ConstraintState.QUADRATIC - wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref) - # middle zone - else: - dm = math.safe_div(efc_D_in[worldid, efcid0], mu * mu * (1.0 + mu * mu)) - nmt = N - mu * T - - force = -dm * nmt * mu - - if efcid == efcid0: - efc_force_out[worldid, efcid] = force - wp.atomic_add(ctx_cost_out, worldid, 0.5 * dm * nmt * nmt) - else: - efc_force_out[worldid, efcid] = -math.safe_div(force, T) * ufrictionj - - efc_state_out[worldid, efcid] = types.ConstraintState.CONE + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + for i in range(rownnz): + sparseid = rowadr + i + colind = efc_J_colind_in[worldid, 0, sparseid] + efc_J = efc_J_in[worldid, 0, sparseid] + wp.atomic_add(qfrc_constraint_out[worldid], colind, efc_J * force) @wp.kernel -def update_constraint_init_qfrc_constraint( - # Data in: - nefc_in: wp.array(dtype=int), - efc_J_in: wp.array3d(dtype=float), - efc_force_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - ctx_done_in: wp.array(dtype=bool), - # Data out: - qfrc_constraint_out: wp.array2d(dtype=float), +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), + njmax_in: int, + # In: + ctx_done_in: wp.array(dtype=bool), + # Data out: + qfrc_constraint_out: wp.array2d(dtype=float), ): worldid, dofid = wp.tid() @@ -1937,7 +2082,112 @@ def update_constraint_gauss_cost(nv: int, dofs_per_thread: int): return kernel -def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext): +@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), + # In: + changed_ids_in: wp.array2d(dtype=int), + changed_count_in: wp.array(dtype=int), + # Out: + ctx_h_out: wp.array3d(dtype=float), +): + """Incrementally update lower triangle of H for changed constraints. + + Each thread handles one (i, j) element of the lower triangle. + For each changed constraint, adds or subtracts D * J[i] * J[j]. + """ + worldid, elementid = wp.tid() + + n_changes = changed_count_in[worldid] + if n_changes == 0: + return + + # Lower triangle index: elementid -> (i, j) where i >= j + i = (int(wp.sqrt(float(1 + 8 * elementid))) - 1) // 2 + j = elementid - (i * (i + 1)) // 2 + + delta = float(0.0) + for change_idx in range(n_changes): + efcid = changed_ids_in[worldid, change_idx] + Ji = efc_J_in[worldid, efcid, i] + if Ji == 0.0: + continue + Jj = efc_J_in[worldid, efcid, j] + if Jj == 0.0: + continue + + D = efc_D_in[worldid, efcid] + if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value: + delta += D * Ji * Jj + else: + delta -= D * Ji * Jj + + if delta != 0.0: + ctx_h_out[worldid, i, j] += delta + + +@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), + # In: + changed_ids_in: wp.array2d(dtype=int), + changed_count_in: wp.array(dtype=int), + # Out: + ctx_h_out: wp.array3d(dtype=float), +): + """Incrementally update lower triangle of H for changed constraints (sparse J).""" + worldid, change_idx = wp.tid() + + n_changes = changed_count_in[worldid] + if change_idx >= n_changes: + return + + efcid = changed_ids_in[worldid, change_idx] + D = efc_D_in[worldid, efcid] + sign = float(0.0) + if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value: + sign = D + else: + sign = -D + + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + + for ii in range(rownnz): + sparseidi = rowadr + ii + Ji = efc_J_in[worldid, 0, sparseidi] + if Ji == 0.0: + continue + colindi = efc_J_colind_in[worldid, 0, sparseidi] + for jj in range(ii + 1): + sparseidj = rowadr + jj + Jj = efc_J_in[worldid, 0, sparseidj] + if Jj == 0.0: + continue + colindj = efc_J_colind_in[worldid, 0, sparseidj] + h = sign * Ji * Jj + # Ensure lower triangle: larger index first + if colindi >= colindj: + wp.atomic_add(ctx_h_out[worldid, colindi], colindj, h) + else: + wp.atomic_add(ctx_h_out[worldid, colindj], colindi, h) + + +def _update_constraint( + m: types.Model, + d: types.Data, + ctx: SolverContext | InverseContext, + track_changes: bool = False, +): """Update constraint arrays after each solve iteration.""" wp.launch( update_constraint_init_cost, @@ -1946,10 +2196,7 @@ def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | Inver outputs=[ctx.gauss, ctx.cost, ctx.prev_cost], ) - wp.launch( - update_constraint_efc, - dim=(d.nworld, d.njmax), - inputs=[ + efc_inputs = [ m.opt.impratio_invsqrt, d.ne, d.nf, @@ -1964,17 +2211,45 @@ def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | Inver d.nacon, ctx.Jaref, ctx.done, - ], - outputs=[d.efc.force, d.efc.state, ctx.cost], + ] + + wp.launch( + update_constraint_efc(track_changes), + dim=(d.nworld, d.njmax), + inputs=efc_inputs, + outputs=[ + d.efc.force, + d.efc.state, + ctx.cost, + ctx.changed_efc_ids, + ctx.changed_efc_count, + ], ) # qfrc_constraint = efc_J.T @ efc_force - wp.launch( - update_constraint_init_qfrc_constraint, - dim=(d.nworld, m.nv), - inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, ctx.done], - outputs=[d.qfrc_constraint], - ) + if SPARSE_CONSTRAINT_JACOBIAN: + d.qfrc_constraint.zero_() + wp.launch( + update_constraint_init_qfrc_constraint_sparse, + dim=(d.nworld, d.njmax), + inputs=[ + d.nefc, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.force, + ctx.done, + ], + outputs=[d.qfrc_constraint], + ) + else: + wp.launch( + update_constraint_init_qfrc_constraint_dense, + dim=(d.nworld, m.nv), + inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, ctx.done], + outputs=[d.qfrc_constraint], + ) # if we are only using 1 thread, it makes sense to do more dofs and skip the atomics. # For more than 1 thread, dofs_per_thread is lower for better load balancing. @@ -2140,7 +2415,7 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int): @cache_kernel -def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int): +def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int): if njmax < tile_size: tile_size = njmax @@ -2166,7 +2441,9 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int nefc = nefc_in[worldid] - sum_val = wp.tile_load(qM_in[worldid], shape=(nv_padded, nv_padded), bounds_check=True) + sum_val = wp.tile_load( + qM_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True + ) # Each tile processes one output tile by looping over all constraints for k in range(0, njmax, TILE_SIZE_K): @@ -2176,7 +2453,12 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int # AD: leaving bounds-check disabled here because I'm not entirely sure that # everything always hits the fast path. The padding takes care of any # potential OOB accesses. - J_kj = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv_padded), offset=(k, 0), bounds_check=False) + J_kj = wp.tile_load( + efc_J_in[worldid], + shape=(TILE_SIZE_K, nv_pad), + offset=(k, 0), + bounds_check=False, + ) # state check D_k = wp.tile_load(efc_D_in[worldid], shape=TILE_SIZE_K, offset=k, bounds_check=False) @@ -2191,7 +2473,11 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int active_tile = wp.tile_map(active_check, tid_tile, threshold_tile) D_k = wp.tile_map(wp.mul, active_tile, D_k) - J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_kj), wp.tile_broadcast(D_k, shape=(nv_padded, TILE_SIZE_K))) + J_ki = wp.tile_map( + wp.mul, + wp.tile_transpose(J_kj), + wp.tile_broadcast(D_k, shape=(nv_pad, TILE_SIZE_K)), + ) sum_val += wp.tile_matmul(J_ki, J_kj) @@ -2202,30 +2488,198 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int # TODO(thowell): combine with JTDAJ ? @wp.kernel -def update_gradient_JTCJ( - # Model: - opt_impratio_invsqrt: wp.array(dtype=float), - dof_tri_row: wp.array(dtype=int), - dof_tri_col: wp.array(dtype=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), - naconmax_in: int, - nacon_in: wp.array(dtype=int), - # In: - ctx_Jaref_in: wp.array2d(dtype=float), - ctx_done_in: wp.array(dtype=bool), - nblocks_perblock: int, - dim_block: int, - # Out: - ctx_h_out: wp.array3d(dtype=float), +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), + # 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), + naconmax_in: int, + nacon_in: wp.array(dtype=int), + # In: + ctx_Jaref_in: wp.array2d(dtype=float), + ctx_done_in: wp.array(dtype=bool), + nblocks_perblock: int, + dim_block: int, + # Out: + h_out: wp.array3d(dtype=float), +): + conid_start, elementid = wp.tid() + + dof1id = dof_tri_row[elementid] + dof2id = dof_tri_col[elementid] + + for i in range(nblocks_perblock): + conid = conid_start + i * dim_block + + if conid >= min(nacon_in[0], naconmax_in): + return + + worldid = contact_worldid_in[conid] + if ctx_done_in[worldid]: + return + + condim = contact_dim_in[conid] + + if condim == 1: + return + + # check contact status + if contact_dist_in[conid] - contact_includemargin_in[conid] >= 0.0: + return + + efcid0 = contact_efc_address_in[conid, 0] + if efc_state_in[worldid, efcid0] != types.ConstraintState.CONE: + return + + fri = contact_friction_in[conid] + mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] + + mu2 = mu * mu + dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2)) + + if dm == 0.0: + return + + n = ctx_Jaref_in[worldid, efcid0] * mu + u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0) + + tt = float(0.0) + for j in range(1, condim): + efcidj = contact_efc_address_in[conid, j] + uj = ctx_Jaref_in[worldid, efcidj] * fri[j - 1] + tt += uj * uj + u[j] = uj + + if tt <= 0.0: + t = 0.0 + else: + t = wp.sqrt(tt) + t = wp.max(t, types.MJ_MINVAL) + ttt = wp.max(t * t * t, types.MJ_MINVAL) + + h = float(0.0) + + for dim1id in range(condim): + if dim1id == 0: + efcid1 = efcid0 + else: + efcid1 = contact_efc_address_in[conid, dim1id] + + # TODO(team): improve performance for sparse code path + rownnz1 = efc_J_rownnz_in[worldid, efcid1] + rowadr1 = efc_J_rowadr_in[worldid, efcid1] + + efc_J11 = float(0.0) + efc_J12 = float(0.0) + for i1 in range(rownnz1): + sparseid1 = rowadr1 + i1 + colind1 = efc_J_colind_in[worldid, 0, sparseid1] + if dof1id == colind1: + efc_J11 = efc_J_in[worldid, 0, sparseid1] + if dof2id == colind1: + efc_J12 = efc_J_in[worldid, 0, sparseid1] + if efc_J11 != 0.0 and efc_J12 != 0.0: + break + + ui = u[dim1id] + + for dim2id in range(0, dim1id + 1): + if dim2id == 0: + efcid2 = efcid0 + else: + efcid2 = contact_efc_address_in[conid, dim2id] + + rownnz2 = efc_J_rownnz_in[worldid, efcid2] + rowadr2 = efc_J_rowadr_in[worldid, efcid2] + + efc_J21 = float(0.0) + efc_J22 = float(0.0) + for i2 in range(rownnz2): + sparseid2 = rowadr2 + i2 + colind2 = efc_J_colind_in[worldid, 0, sparseid2] + if dof1id == colind2: + efc_J21 = efc_J_in[worldid, 0, sparseid2] + if dof2id == colind2: + efc_J22 = efc_J_in[worldid, 0, sparseid2] + if efc_J21 != 0.0 and efc_J22 != 0.0: + break + + uj = u[dim2id] + + # set first row/column: (1, -mu/t * u) + if dim1id == 0 and dim2id == 0: + hcone = 1.0 + elif dim1id == 0: + hcone = -math.safe_div(mu, t) * uj + elif dim2id == 0: + hcone = -math.safe_div(mu, t) * ui + else: + hcone = mu * math.safe_div(n, ttt) * ui * uj + + # add to diagonal: mu^2 - mu * n / t + if dim1id == dim2id: + hcone += mu2 - mu * math.safe_div(n, t) + + # pre and post multiply by diag(mu, friction) scale by dm + if dim1id == 0: + fri1 = mu + else: + fri1 = fri[dim1id - 1] + + if dim2id == 0: + fri2 = mu + else: + fri2 = fri[dim2id - 1] + + hcone *= dm * fri1 * fri2 + + if hcone != 0.0: + h += hcone * efc_J11 * efc_J22 + + if dim1id != dim2id: + h += hcone * efc_J12 * efc_J21 + + h_out[worldid, dof1id, dof2id] += h + + +@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), + # 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), + naconmax_in: int, + nacon_in: wp.array(dtype=int), + # In: + ctx_Jaref_in: wp.array2d(dtype=float), + ctx_done_in: wp.array(dtype=bool), + nblocks_perblock: int, + dim_block: int, + # Out: + ctx_h_out: wp.array3d(dtype=float), ): conid_start, elementid = wp.tid() @@ -2256,7 +2710,7 @@ def update_gradient_JTCJ( continue fri = contact_friction_in[conid] - mu = fri[0] * opt_impratio_invsqrt[worldid] + mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] mu2 = mu * mu dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2)) @@ -2409,6 +2863,40 @@ def padding_h(nv: int, ctx_done_in: wp.array(dtype=bool), ctx_h_out: wp.array3d( ctx_h_out[worldid, dofid, dofid] = 1.0 +def _cholesky_factorize_solve( + m: types.Model, d: types.Data, ctx: SolverContext +): + """Cholesky factorize ctx.h and solve for Mgrad.""" + if m.nv <= _BLOCK_CHOLESKY_DIM: + wp.launch_tiled( + update_gradient_cholesky(m.nv), + dim=d.nworld, + inputs=[ctx.grad, ctx.h, ctx.done], + outputs=[ctx.Mgrad], + block_dim=m.block_dim.update_gradient_cholesky, + ) + else: + wp.launch( + padding_h, + dim=(d.nworld, m.nv_pad - m.nv), + inputs=[m.nv, ctx.done], + outputs=[ctx.h], + ) + + wp.launch_tiled( + update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad), + dim=d.nworld, + inputs=[ + ctx.done, + ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), + ctx.h, + ctx.hfactor, + ], + outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))], + block_dim=m.block_dim.update_gradient_cholesky_blocked, + ) + + def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): # grad = Ma - qfrc_smooth - qfrc_constraint wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.grad_dot]) @@ -2424,22 +2912,125 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): smooth.solve_m(m, d, ctx.Mgrad, ctx.grad) elif m.opt.solver == types.SolverType.NEWTON: # h = qM + (efc_J.T * efc_D * active) @ efc_J - if m.is_sparse: + if SPARSE_CONSTRAINT_JACOBIAN: + # TODO(team): improve performance for sparse code path + @wp.kernel(module="unique", enable_backward=False) + 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), + # In: + ctx_done_in: wp.array(dtype=bool), + # Out: + h_out: wp.array3d(dtype=float), + ): + worldid, efcid = wp.tid() + + if ctx_done_in[worldid]: + return + + if efcid >= nefc_in[worldid]: + return + + efc_D = efc_D_in[worldid, efcid] + efc_state = efc_state_in[worldid, efcid] + + if state_check(efc_D, efc_state) == 0.0: + return + + rownnz = efc_J_rownnz_in[worldid, efcid] + rowadr = efc_J_rowadr_in[worldid, efcid] + + for i in range(rownnz): + sparseidi = rowadr + i + Ji = efc_J_in[worldid, 0, sparseidi] + colindi = efc_J_colind_in[worldid, 0, sparseidi] + for j in range(i, rownnz): + if j == i: + sparseidj = sparseidi + Jj = Ji + colindj = colindi + else: + sparseidj = rowadr + j + Jj = efc_J_in[worldid, 0, sparseidj] + 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) + + wp.launch( + _JTDAJ_sparse, + dim=(d.nworld, d.njmax), + inputs=[ + d.nefc, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.D, + d.efc.state, + ctx.done, + ], + outputs=[ctx.h], + ) + + if m.is_sparse: + wp.launch( + update_gradient_set_h_qM_lower_sparse, + dim=(d.nworld, m.qM_fullm_i.size), + inputs=[m.qM_fullm_i, m.qM_fullm_j, d.qM, ctx.done], + outputs=[ctx.h], + ) + else: + # dense M: copy qM directly into h + @wp.kernel(module="unique", enable_backward=False) + def _set_h_qM_dense( + nv: int, + qM_in: wp.array3d(dtype=float), + ctx_done_in: wp.array(dtype=bool), + ctx_h_out: wp.array3d(dtype=float), + ): + worldid, i, j = wp.tid() + if ctx_done_in[worldid]: + return + if i >= nv or j >= nv: + return + if i >= j: + ctx_h_out[worldid, i, j] += qM_in[worldid, i, j] + + wp.launch( + _set_h_qM_dense, + dim=(d.nworld, m.nv, m.nv), + inputs=[m.nv, d.qM, ctx.done], + outputs=[ctx.h], + ) + elif m.is_sparse: num_blocks_ceil = ceil(m.nv / types.TILE_SIZE_JTDAJ_SPARSE) lower_triangle_dim = int(num_blocks_ceil * (num_blocks_ceil + 1) / 2) - wp.launch_tiled( - update_gradient_JTDAJ_sparse_tiled(types.TILE_SIZE_JTDAJ_SPARSE, d.njmax), - dim=(d.nworld, lower_triangle_dim), - inputs=[ - d.nefc, - d.efc.J, - d.efc.D, - d.efc.state, - ctx.done, - ], - outputs=[ctx.h], - block_dim=m.block_dim.update_gradient_JTDAJ_sparse, - ) + with scoped_mathdx_gemm_disabled(): + wp.launch_tiled( + update_gradient_JTDAJ_sparse_tiled( + types.TILE_SIZE_JTDAJ_SPARSE, d.njmax + ), + dim=(d.nworld, lower_triangle_dim), + inputs=[ + d.nefc, + d.efc.J, + d.efc.D, + d.efc.state, + ctx.done, + ], + outputs=[ctx.h], + block_dim=m.block_dim.update_gradient_JTDAJ_sparse, + ) wp.launch( update_gradient_set_h_qM_lower_sparse, @@ -2448,21 +3039,23 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): outputs=[ctx.h], ) else: - nv_padded = d.efc.J.shape[2] - wp.launch_tiled( - update_gradient_JTDAJ_dense_tiled(nv_padded, types.TILE_SIZE_JTDAJ_DENSE, d.njmax), - dim=d.nworld, - inputs=[ - d.nefc, - d.qM, - d.efc.J, - d.efc.D, - d.efc.state, - ctx.done, - ], - outputs=[ctx.h], - block_dim=m.block_dim.update_gradient_JTDAJ_dense, - ) + with scoped_mathdx_gemm_disabled(): + wp.launch_tiled( + update_gradient_JTDAJ_dense_tiled( + m.nv_pad, types.TILE_SIZE_JTDAJ_DENSE, d.njmax + ), + dim=d.nworld, + inputs=[ + d.nefc, + d.qM, + d.efc.J, + d.efc.D, + d.efc.state, + ctx.done, + ], + outputs=[ctx.h], + block_dim=m.block_dim.update_gradient_JTDAJ_dense, + ) if m.opt.cone == types.ConeType.ELLIPTIC: # Optimization: launching update_gradient_JTCJ with limited number of blocks on a GPU. @@ -2488,59 +3081,124 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): nblocks_perblock = int((d.naconmax + dim_block - 1) / dim_block) - wp.launch( - update_gradient_JTCJ, - dim=(dim_block, m.dof_tri_row.size), - inputs=[ - m.opt.impratio_invsqrt, - m.dof_tri_row, - m.dof_tri_col, - d.contact.dist, - d.contact.includemargin, - d.contact.friction, - d.contact.dim, - d.contact.efc_address, - d.contact.worldid, - d.efc.J, - d.efc.D, - d.efc.state, - d.naconmax, - d.nacon, - ctx.Jaref, - ctx.done, - nblocks_perblock, - dim_block, - ], - outputs=[ctx.h], - ) + if SPARSE_CONSTRAINT_JACOBIAN: + wp.launch( + update_gradient_JTCJ_sparse, + dim=(d.naconmax, m.dof_tri_row.size), + inputs=[ + m.opt.impratio_invsqrt, + m.dof_tri_row, + m.dof_tri_col, + d.contact.dist, + d.contact.includemargin, + d.contact.friction, + d.contact.dim, + d.contact.efc_address, + d.contact.worldid, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.D, + d.efc.state, + d.naconmax, + d.nacon, + ctx.Jaref, + ctx.done, + nblocks_perblock, + dim_block, + ], + outputs=[ctx.h], + ) + else: + wp.launch( + update_gradient_JTCJ_dense, + dim=(dim_block, m.dof_tri_row.size), + inputs=[ + m.opt.impratio_invsqrt, + m.dof_tri_row, + m.dof_tri_col, + d.contact.dist, + d.contact.includemargin, + d.contact.friction, + d.contact.dim, + d.contact.efc_address, + d.contact.worldid, + d.efc.J, + d.efc.D, + d.efc.state, + d.naconmax, + d.nacon, + ctx.Jaref, + ctx.done, + nblocks_perblock, + dim_block, + ], + outputs=[ctx.h], + ) - if m.nv <= _BLOCK_CHOLESKY_DIM: - wp.launch_tiled( - update_gradient_cholesky(m.nv), - dim=d.nworld, - inputs=[ctx.grad, ctx.h, ctx.done], - outputs=[ctx.Mgrad], - block_dim=m.block_dim.update_gradient_cholesky, - ) - else: - wp.launch( - padding_h, - dim=(d.nworld, m.nv_pad - m.nv), - inputs=[m.nv, ctx.done], - outputs=[ctx.h], - ) - - wp.launch_tiled( - update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad), - dim=d.nworld, - inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.hfactor], - outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))], - block_dim=m.block_dim.update_gradient_cholesky_blocked, - ) + _cholesky_factorize_solve(m, d, ctx) else: raise ValueError(f"Unknown solver type: {m.opt.solver}") +def _update_gradient_incremental( + m: types.Model, d: types.Data, ctx: SolverContext +): + """Incremental gradient update: update H for changed constraints + re-factorize. + + Skips the full J^T*D*J rebuild by applying only the delta from constraints + that changed QUADRATIC state, then re-factorizes and solves. + """ + wp.launch( + update_gradient_zero_grad_dot, + dim=(d.nworld), + inputs=[ctx.done], + outputs=[ctx.grad_dot], + ) + + wp.launch( + update_gradient_grad, + dim=(d.nworld, m.nv), + inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done], + outputs=[ctx.grad, ctx.grad_dot], + ) + + # Update lower triangle of H with delta from changed constraints + if SPARSE_CONSTRAINT_JACOBIAN: + wp.launch( + update_gradient_h_incremental_sparse, + dim=(d.nworld, ctx.changed_efc_ids.shape[1]), + inputs=[ + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.D, + d.efc.state, + ctx.changed_efc_ids, + ctx.changed_efc_count, + ], + outputs=[ctx.h], + ) + else: + lower_tri_dim = m.nv * (m.nv + 1) // 2 + wp.launch( + update_gradient_h_incremental, + dim=(d.nworld, lower_tri_dim), + inputs=[ + d.efc.J, + d.efc.D, + d.efc.state, + ctx.changed_efc_ids, + ctx.changed_efc_count, + ], + outputs=[ctx.h], + ) + + _cholesky_factorize_solve(m, d, ctx) + + @wp.kernel def solve_prev_grad_Mgrad( # In: @@ -2685,8 +3343,25 @@ def _solver_iteration( outputs=[ctx.prev_grad, ctx.prev_Mgrad], ) - _update_constraint(m, d, ctx) - _update_gradient(m, d, ctx) + # Incremental H is only valid for non-elliptic cones. The elliptic cone + # path in update_constraint_efc has early returns that skip state change + # tracking, and the additional JTCJ Hessian term depends on Jaref which + # changes every iteration. + incremental = ( + m.opt.solver == types.SolverType.NEWTON + and m.opt.cone != types.ConeType.ELLIPTIC + ) + + if incremental: + # Must complete before update_constraint_efc which atomically increments. + ctx.changed_efc_count.zero_() + + _update_constraint(m, d, ctx, track_changes=incremental) + + if incremental: + _update_gradient_incremental(m, d, ctx) + else: + _update_gradient(m, d, ctx) # polak-ribiere if m.opt.solver == types.SolverType.CG: @@ -2747,10 +3422,18 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont ctx.Jaref.zero_() wp.launch( - solve_init_jaref(m.nv, dofs_per_thread), - dim=(d.nworld, d.njmax, threads_per_efc), - inputs=[d.nefc, d.qacc, d.efc.J, d.efc.aref], - outputs=[ctx.Jaref], + solve_init_jaref(m.is_sparse, m.nv, dofs_per_thread), + dim=(d.nworld, d.njmax, threads_per_efc), + inputs=[ + d.nefc, + d.qacc, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, + d.efc.J, + d.efc.aref, + ], + outputs=[ctx.Jaref], ) # Ma = qM @ qacc 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 bc22a82b..94434ff9 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -14,6 +14,7 @@ # ============================================================================== import dataclasses import enum +from typing import Callable import mujoco import warp as wp @@ -32,6 +33,9 @@ MJ_MAX_EPAFACES = 5 TILE_SIZE_JTDAJ_SPARSE = 16 TILE_SIZE_JTDAJ_DENSE = 16 +# TODO(team): remove after improving performance for sparse constraint jacobian +SPARSE_CONSTRAINT_JACOBIAN = False + # TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml TEXTURE_DTYPE = wp.Texture2D if hasattr(wp, "Texture2D") else int @@ -134,6 +138,20 @@ class ProjectionType(enum.IntEnum): ORTHOGRAPHIC = 1 +class Stage(enum.IntEnum): + """Computation stage. + + Attributes: + POS: position-dependent + VEL: velocity-dependent + ACC: acceleration/force-dependent + """ + + POS = mujoco.mjtStage.mjSTAGE_POS + VEL = mujoco.mjtStage.mjSTAGE_VEL + ACC = mujoco.mjtStage.mjSTAGE_ACC + + class DataType(enum.IntFlag): """Sensor data types. @@ -167,6 +185,7 @@ class DisableBit(enum.IntFlag): SENSOR: sensors EULERDAMP: implicit damping for Euler integration NATIVECCD: native convex collision detection (ignored in MJWarp) + ISLAND: constraint islands """ CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT @@ -185,7 +204,8 @@ class DisableBit(enum.IntFlag): SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP NATIVECCD = mujoco.mjtDisableBit.mjDSBL_NATIVECCD - # unsupported: MIDPHASE, AUTORESET, ISLAND + ISLAND = mujoco.mjtDisableBit.mjDSBL_ISLAND + # unsupported: MIDPHASE, AUTORESET class EnableBit(enum.IntFlag): @@ -232,6 +252,7 @@ class DynType(enum.IntEnum): FILTER: linear filter: da/dt = (u-a) / tau FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration MUSCLE: piece-wise linear filter with two time constants + USER: user-defined dynamics via act_dyn_callback """ NONE = mujoco.mjtDyn.mjDYN_NONE @@ -239,7 +260,7 @@ class DynType(enum.IntEnum): FILTER = mujoco.mjtDyn.mjDYN_FILTER FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE - # unsupported: USER + USER = mujoco.mjtDyn.mjDYN_USER class GainType(enum.IntEnum): @@ -249,12 +270,13 @@ class GainType(enum.IntEnum): FIXED: fixed gain AFFINE: const + kp*length + kv*velocity MUSCLE: muscle FLV curve computed by muscle_gain + USER: user-defined gain via act_gain_callback """ FIXED = mujoco.mjtGain.mjGAIN_FIXED AFFINE = mujoco.mjtGain.mjGAIN_AFFINE MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE - # unsupported: USER + USER = mujoco.mjtGain.mjGAIN_USER class BiasType(enum.IntEnum): @@ -264,12 +286,13 @@ class BiasType(enum.IntEnum): NONE: no bias AFFINE: const + kp*length + kv*velocity MUSCLE: muscle passive force computed by muscle_bias + USER: user-defined bias via act_bias_callback """ NONE = mujoco.mjtBias.mjBIAS_NONE AFFINE = mujoco.mjtBias.mjBIAS_AFFINE MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE - # unsupported: USER + USER = mujoco.mjtBias.mjBIAS_USER class JointType(enum.IntEnum): @@ -462,6 +485,7 @@ class SensorType(enum.IntEnum): FRAMELINACC: 3D linear acceleration FRAMEANGACC: 3D angular acceleration TACTILE: tactile sensor + USER: user-defined sensor via sensor_callback """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER @@ -511,6 +535,7 @@ class SensorType(enum.IntEnum): FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC TACTILE = mujoco.mjtSensor.mjSENS_TACTILE + USER = mujoco.mjtSensor.mjSENS_USER class ObjType(enum.IntEnum): @@ -768,6 +793,29 @@ class TileSet: size: int +@dataclasses.dataclass +class Callback: + """Callbacks for custom physics behavior. + + Attributes: + passive: custom passive forces, writes to ``Data.qfrc_passive`` + control: custom control laws, writes to ``Data.ctrl`` + act_dyn: custom actuator dynamics, writes to ``Data.act_dot`` + act_gain: custom actuator gains, writes to ``Data.actuator_force`` + act_bias: custom actuator biases, writes to ``Data.actuator_force`` + sensor: custom sensors, writes to ``Data.sensordata`` + contactfilter: custom contact filtering, writes to ``Data.contact`` + """ + + passive: Callable | None = None + control: Callable | None = None + act_dyn: Callable | None = None + act_gain: Callable | None = None + act_bias: Callable | None = None + sensor: Callable | None = None + contactfilter: Callable | None = None + + @dataclasses.dataclass class Model: """Model definition and parameters. @@ -793,6 +841,7 @@ class Model: nflexelem: number of elements in all flexes nflexelemdata: number of element vertex ids in all flexes nflexelemedge: number of element edge ids in all flexes + nJfe: number of non-zeros in sparse flexedge Jacobian nmesh: number of meshes nmeshvert: number of vertices for all meshes nmeshnormal: number of normals in all meshes @@ -812,6 +861,7 @@ class Model: nsensor: number of sensors nmocap: number of mocap bodies nplugin: number of plugin instances + nJmom: number of non-zeros in actuator_moment ngravcomp: number of bodies with nonzero gravcomp nsensordata: number of elements in sensor data vector opt: physics options @@ -947,7 +997,9 @@ class Model: mesh_vertadr: first vertex address (nmesh,) mesh_vertnum: number of vertices (nmesh,) mesh_faceadr: first face address (nmesh,) + mesh_octadr: octree address for each mesh (nmesh,) mesh_normaladr: first normal address (nmesh,) + mesh_normalnum: number of normals (nmesh,) mesh_graphadr: graph data address; -1: no graph (nmesh,) mesh_vert: vertex positions for all meshes (nmeshvert, 3) mesh_normal: normals for all meshes (nmeshnormal, 3) @@ -968,7 +1020,8 @@ class Model: hfield_ncol: number of columns in grid (nhfield,) hfield_adr: start address in hfield_data (nhfield,) hfield_data: elevation data (nhfielddata,) - mat_texid: texture id for rendering (*, nmat, mjNTEXROLE) + mat_texid: texture id for rendering (*, nmat, + mjNTEXROLE) mat_texrepeat: texture repeat for rendering (*, nmat, 2) mat_rgba: rgba (*, nmat, 4) pair_dim: contact dimensionality (npair,) @@ -993,10 +1046,14 @@ class Model: tendon_num: number of objects in tendon's path (ntendon,) tendon_limited: does tendon have length limits (ntendon,) tendon_actfrclimited: does ten have actuator force limit (ntendon,) - tendon_solref_lim: constraint solver reference: limit (*, ntendon, mjNREF) - tendon_solimp_lim: constraint solver impedance: limit (*, ntendon, mjNIMP) - tendon_solref_fri: constraint solver reference: friction (*, ntendon, mjNREF) - tendon_solimp_fri: constraint solver impedance: friction (*, ntendon, mjNIMP) + tendon_solref_lim: constraint solver reference: limit (*, ntendon, + mjNREF) + tendon_solimp_lim: constraint solver impedance: limit (*, ntendon, + mjNIMP) + tendon_solref_fri: constraint solver reference: friction (*, ntendon, + mjNREF) + tendon_solimp_fri: constraint solver impedance: friction (*, ntendon, + mjNIMP) tendon_range: tendon length limits (*, ntendon, 2) tendon_actfrcrange: range of total actuator force (*, ntendon, 2) tendon_margin: min distance for limit detection (*, ntendon) @@ -1028,9 +1085,9 @@ class Model: actuator_forcerange: range of forces (*, nu, 2) actuator_actrange: range of activations (*, nu, 2) actuator_gear: scale length and transmitted force (*, nu, 6) - actuator_cranklength: crank length for slider-crank (nu,) - actuator_acc0: acceleration from unit force in qpos0 (nu,) - actuator_lengthrange: feasible actuator length range (nu, 2) + actuator_cranklength: crank length for slider-crank (*, nu) + actuator_acc0: acceleration from unit force in qpos0 (*, nu) + actuator_lengthrange: feasible actuator length range (*, nu, 2) sensor_type: sensor type (SensorType) (nsensor,) sensor_datatype: numeric data type (DataType) (nsensor,) sensor_objtype: type of sensorized object (ObjType) (nsensor,) @@ -1049,6 +1106,7 @@ class Model: mapM2M: index mapping from M (legacy) to M (CSR) (nC) warp only fields: + callback: custom physics callbacks nbranch: number of branches (leaf-to-root paths) nv_pad: number of degrees of freedom + padding nacttrnbody: number of actuators with body transmission @@ -1062,12 +1120,14 @@ class Model: nmaxpolygon: maximum number of verts per polygon nmaxmeshdeg: maximum number of polygons per vert is_sparse: whether to use sparse representations - has_fluid: True if wind, density, or viscosity are non-zero at put_model time + has_fluid: True if wind, density, or viscosity are non-zero at put_model + time has_sdf_geom: whether the model contains SDF geoms block_dim: block dim options body_tree: list of body ids by tree level body_branches: flattened body ids for all branches - body_branch_start: start index in body_branches for each branch (nbranch + 1,) + body_branch_start: start index in body_branches for each branch (nbranch + + 1,) mocap_bodyid: id of body for mocap (nmocap,) body_fluid_ellipsoid: does body use ellipsoid fluid (nbody,) jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr @@ -1156,6 +1216,7 @@ class Model: nflexelem: int nflexelemdata: int nflexelemedge: int + nJfe: int nmesh: int nmeshvert: int nmeshnormal: int @@ -1175,6 +1236,7 @@ class Model: nsensor: int nmocap: int nplugin: int + nJmom: int ngravcomp: int nsensordata: int opt: Option @@ -1306,11 +1368,13 @@ class Model: flex_damping: array("nflex", float) flexedge_J_rownnz: array("nflexedge", int) flexedge_J_rowadr: array("nflexedge", int) - flexedge_J_colind: wp.array(dtype=int) + flexedge_J_colind: array("nJfe", int) mesh_vertadr: array("nmesh", int) mesh_vertnum: array("nmesh", int) mesh_faceadr: array("nmesh", int) + mesh_octadr: array("nmesh", int) mesh_normaladr: array("nmesh", int) + mesh_normalnum: array("nmesh", int) mesh_graphadr: array("nmesh", int) mesh_vert: array("nmeshvert", wp.vec3) mesh_normal: array("nmeshnormal", wp.vec3) @@ -1391,9 +1455,9 @@ class Model: actuator_forcerange: array("*", "nu", wp.vec2) actuator_actrange: array("*", "nu", wp.vec2) actuator_gear: array("*", "nu", wp.spatial_vector) - actuator_cranklength: array("nu", float) - actuator_acc0: array("nu", float) - actuator_lengthrange: array("nu", wp.vec2) + actuator_cranklength: array("*", "nu", float) + actuator_acc0: array("*", "nu", float) + actuator_lengthrange: array("*", "nu", wp.vec2) sensor_type: array("nsensor", int) sensor_datatype: array("nsensor", int) sensor_objtype: array("nsensor", int) @@ -1411,6 +1475,7 @@ class Model: M_colind: array("nC", int) mapM2M: array("nC", int) # warp only fields: + callback: Callback nbranch: int nv_pad: int nacttrnbody: int @@ -1546,7 +1611,14 @@ class Constraint: Attributes: type: constraint type (ConstraintType) (nworld, njmax) id: id of object of specific type (nworld, njmax) - J: constraint Jacobian (nworld, njmax_pad, nv_pad) + J_rownnz: number of non-zeros in J row (nworld, 0) dense (nworld, + njmax) sparse + J_rowadr: row start address in colind array (nworld, 0) dense (nworld, + njmax) sparse + J_colind: column indices in J (nworld, 0, 0) dense + (nworld, 1, njmax * nv) sparse + J: constraint Jacobian (nworld, njmax_pad, + nv_pad) dense (nworld, 1, njmax * nv) sparse pos: constraint position (equality, contact) (nworld, njmax) margin: inclusion margin (contact) (nworld, njmax) D: constraint mass (nworld, njmax_pad) @@ -1555,13 +1627,17 @@ class Constraint: frictionloss: frictionloss (friction) (nworld, njmax) force: constraint force in constraint space (nworld, njmax) state: constraint state (nworld, njmax_pad) + warp only fields: Ma: M*qacc (nworld, nv) """ type: array("nworld", "njmax", int) id: array("nworld", "njmax", int) - J: array("nworld", "njmax_pad", "nv_pad", float) + J_rownnz: wp.array2d(dtype=int) + J_rowadr: wp.array2d(dtype=int) + J_colind: wp.array3d(dtype=int) + J: wp.array3d(dtype=float) pos: array("nworld", "njmax", float) margin: array("nworld", "njmax", float) D: array("nworld", "njmax_pad", float) @@ -1583,6 +1659,7 @@ class Data: nf: number of friction constraints (nworld,) nl: number of limit constraints (nworld,) nefc: number of constraints (nworld,) + nisland: number of constraint islands (nworld,) time: simulation time (nworld,) energy: potential, kinetic energy (nworld, 2) qpos: position (nworld, nq) @@ -1591,52 +1668,88 @@ class Data: qacc_warmstart: acceleration used for warmstart (nworld, nv) ctrl: control (nworld, nu) qfrc_applied: applied generalized force (nworld, nv) - xfrc_applied: applied Cartesian force/torque (nworld, nbody, 6) + xfrc_applied: applied Cartesian force/torque (nworld, nbody, + 6) eq_active: enable/disable constraints (nworld, neq) - mocap_pos: position of mocap bodies (nworld, nmocap, 3) - mocap_quat: orientation of mocap bodies (nworld, nmocap, 4) + mocap_pos: position of mocap bodies (nworld, nmocap, + 3) + mocap_quat: orientation of mocap bodies (nworld, nmocap, + 4) qacc: acceleration (nworld, nv) act_dot: time-derivative of actuator activation (nworld, na) - sensordata: sensor data array (nworld, nsensordata,) - xpos: Cartesian position of body frame (nworld, nbody, 3) - xquat: Cartesian orientation of body frame (nworld, nbody, 4) - xmat: Cartesian orientation of body frame (nworld, nbody, 3, 3) - xipos: Cartesian position of body com (nworld, nbody, 3) - ximat: Cartesian orientation of body inertia (nworld, nbody, 3, 3) - xanchor: Cartesian position of joint anchor (nworld, njnt, 3) - xaxis: Cartesian joint axis (nworld, njnt, 3) - geom_xpos: Cartesian geom position (nworld, ngeom, 3) - geom_xmat: Cartesian geom orientation (nworld, ngeom, 3, 3) - site_xpos: Cartesian site position (nworld, nsite, 3) - site_xmat: Cartesian site orientation (nworld, nsite, 3, 3) - cam_xpos: Cartesian camera position (nworld, ncam, 3) - cam_xmat: Cartesian camera orientation (nworld, ncam, 3, 3) - light_xpos: Cartesian light position (nworld, nlight, 3) - light_xdir: Cartesian light direction (nworld, nlight, 3) - subtree_com: center of mass of each subtree (nworld, nbody, 3) + sensordata: sensor data array (nworld, + nsensordata,) + xpos: Cartesian position of body frame (nworld, nbody, + 3) + xquat: Cartesian orientation of body frame (nworld, nbody, + 4) + xmat: Cartesian orientation of body frame (nworld, nbody, + 3, 3) + xipos: Cartesian position of body com (nworld, nbody, + 3) + ximat: Cartesian orientation of body inertia (nworld, nbody, + 3, 3) + xanchor: Cartesian position of joint anchor (nworld, njnt, + 3) + xaxis: Cartesian joint axis (nworld, njnt, + 3) + geom_xpos: Cartesian geom position (nworld, ngeom, + 3) + geom_xmat: Cartesian geom orientation (nworld, ngeom, + 3, 3) + site_xpos: Cartesian site position (nworld, nsite, + 3) + site_xmat: Cartesian site orientation (nworld, nsite, + 3, 3) + cam_xpos: Cartesian camera position (nworld, ncam, + 3) + cam_xmat: Cartesian camera orientation (nworld, ncam, + 3, 3) + light_xpos: Cartesian light position (nworld, nlight, + 3) + light_xdir: Cartesian light direction (nworld, nlight, + 3) + subtree_com: center of mass of each subtree (nworld, nbody, + 3) cdof: com-based motion axis of each dof (rot:lin) (nworld, nv, 6) - cinert: com-based body inertia and mass (nworld, nbody, 10) - flexvert_xpos: cartesian flex vertex positions (nworld, nflexvert, 3) - flexedge_J: edge length Jacobian (nworld, 1, nflexedge*6) - flexedge_length: flex edge lengths (nworld, nflexedge, 1) - ten_wrapadr: start address of tendon's path (nworld, ntendon) - ten_wrapnum: number of wrap points in path (nworld, ntendon) - ten_J: tendon Jacobian (nworld, ntendon, nv) - ten_length: tendon lengths (nworld, ntendon) - wrap_obj: geomid; -1: site; -2: pulley (nworld, nwrap, 2) - wrap_xpos: Cartesian 3D points in all paths (nworld, nwrap, 6) + cinert: com-based body inertia and mass (nworld, nbody, + 10) + flexvert_xpos: cartesian flex vertex positions (nworld, + nflexvert, 3) + flexedge_J: edge length Jacobian (nworld, nJfe) + flexedge_length: flex edge lengths (nworld, + nflexedge, 1) + ten_wrapadr: start address of tendon's path (nworld, + ntendon) + ten_wrapnum: number of wrap points in path (nworld, + ntendon) + ten_J: tendon Jacobian (nworld, + ntendon, nv) + ten_length: tendon lengths (nworld, + ntendon) + wrap_obj: geomid; -1: site; -2: pulley (nworld, nwrap, + 2) + wrap_xpos: Cartesian 3D points in all paths (nworld, nwrap, + 6) actuator_length: actuator lengths (nworld, nu) - actuator_moment: actuator moments (nworld, nu, nv) - crb: com-based composite inertia and mass (nworld, nbody, 10) - qM: total inertia (nworld, nv, nv) if dense - (nworld, 1, nM) if sparse - qLD: L'*D*L factorization of M (nworld, nv, nv) if dense - (nworld, 1, nC) if sparse + moment_rownnz: number of non-zeros in actuator_moment row (nworld, nu) + moment_rowadr: row start address in actuator_moment (nworld, nu) + moment_colind: column indices in sparse actuator_moment (nworld, nJmom) + actuator_moment: actuator moments (nworld, nJmom) + crb: com-based composite inertia and mass (nworld, nbody, + 10) + qM: total inertia (nworld, nv, nv) + if dense (nworld, 1, nM) if sparse + qLD: L'*D*L factorization of M (nworld, nv, nv) + if dense (nworld, 1, nC) if sparse qLDiagInv: 1/diag(D) (nworld, nv) - flexedge_velocity: flex edge velocities (nworld, nflexedge) - ten_velocity: tendon velocities (nworld, ntendon) + flexedge_velocity: flex edge velocities (nworld, + nflexedge) + ten_velocity: tendon velocities (nworld, + ntendon) actuator_velocity: actuator velocities (nworld, nu) - cvel: com-based velocity (rot:lin) (nworld, nbody, 6) + cvel: com-based velocity (rot:lin) (nworld, nbody, + 6) cdof_dot: time-derivative of cdof (rot:lin) (nworld, nv, 6) qfrc_bias: C(qpos,qvel) (nworld, nv) qfrc_spring: passive spring force (nworld, nv) @@ -1644,27 +1757,33 @@ class Data: qfrc_gravcomp: passive gravity compensation force (nworld, nv) qfrc_fluid: passive fluid force (nworld, nv) qfrc_passive: total passive force (nworld, nv) - subtree_linvel: linear velocity of subtree com (nworld, nbody, 3) - subtree_angmom: angular momentum about subtree com (nworld, nbody, 3) + subtree_linvel: linear velocity of subtree com (nworld, nbody, + 3) + subtree_angmom: angular momentum about subtree com (nworld, nbody, + 3) actuator_force: actuator force in actuation space (nworld, nu) qfrc_actuator: actuator force (nworld, nv) qfrc_smooth: net unconstrained force (nworld, nv) qacc_smooth: unconstrained acceleration (nworld, nv) qfrc_constraint: constraint force (nworld, nv) qfrc_inverse: net external force; should equal: (nworld, nv) - qfrc_applied + J.T @ xfrc_applied - + qfrc_actuator - cacc: com-based acceleration (nworld, nbody, 6) - cfrc_int: com-based interaction force with parent (nworld, nbody, 6) - cfrc_ext: com-based external force on body (nworld, nbody, 6) + qfrc_applied + J.T @ xfrc_applied + qfrc_actuator + cacc: com-based acceleration (nworld, nbody, + 6) + cfrc_int: com-based interaction force with parent (nworld, nbody, + 6) + cfrc_ext: com-based external force on body (nworld, nbody, + 6) contact: contact data efc: constraint data + tree_island: island ID per tree (-1 if unconstrained) (nworld, ntree) warp only fields: nworld: number of worlds naconmax: maximum number of contacts (shared across all worlds) naccdmax: maximum number of contacts for CCD (all worlds) njmax: maximum number of constraints per world + njmax_pad: njmax rounded up to the nearest multiple of TILE_SIZE_JTDAJ nacon: number of detected contacts (across all worlds) (1,) ncollision: collision count from broadphase (1,) """ @@ -1674,6 +1793,7 @@ class Data: nf: array("nworld", int) nl: array("nworld", int) nefc: array("nworld", int) + nisland: array("nworld", int) time: array("nworld", float) energy: array("nworld", wp.vec2) qpos: array("nworld", "nq", float) @@ -1708,7 +1828,7 @@ class Data: cdof: array("nworld", "nv", wp.spatial_vector) cinert: array("nworld", "nbody", vec10) flexvert_xpos: array("nworld", "nflexvert", wp.vec3) - flexedge_J: wp.array3d(dtype=float) + flexedge_J: array("nworld", "nJfe", float) flexedge_length: array("nworld", "nflexedge", float) ten_wrapadr: array("nworld", "ntendon", int) ten_wrapnum: array("nworld", "ntendon", int) @@ -1717,7 +1837,10 @@ class Data: wrap_obj: array("nworld", "nwrap", wp.vec2i) wrap_xpos: array("nworld", "nwrap", wp.spatial_vector) actuator_length: array("nworld", "nu", float) - actuator_moment: array("nworld", "nu", "nv", float) + moment_rownnz: array("nworld", "nu", int) + moment_rowadr: array("nworld", "nu", int) + 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) @@ -1746,31 +1869,18 @@ class Data: cfrc_ext: array("nworld", "nbody", wp.spatial_vector) contact: Contact efc: Constraint + tree_island: array("nworld", "ntree", int) # warp only fields: nworld: int naconmax: int naccdmax: int njmax: int + njmax_pad: int nacon: array(1, int) ncollision: array(1, int) -@dataclasses.dataclass -class CollisionContext: - """Collision driver intermediate arrays. - - Attributes: - collision_pair: collision pairs from broadphase (naconmax, 2) - collision_pairid: ids from broadphase (naconmax, 2) - collision_worldid: collision world ids from broadphase (naconmax,) - """ - - collision_pair: wp.array - collision_pairid: wp.array - collision_worldid: wp.array - - @dataclasses.dataclass class RenderContext: """Context for rendering. @@ -1794,7 +1904,7 @@ class RenderContext: textures_registry: texture registry hfield_registry: hfield BVH id to warp mesh mapping hfield_bvh_id: hfield BVH ids - hfield_bounds_size: hfield bounds size + hfield_bounds_size: hfield bounds half-extents flex_mesh: flex mesh flex_rgba: flex rgba flex_bvh_id: flex BVH id diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py new file mode 100644 index 00000000..b7debea4 --- /dev/null +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py @@ -0,0 +1,97 @@ +# Copyright 2026 The Newton Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Package version checking utilities.""" + +import importlib.metadata +import operator +import re + + +def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]: + """Parse a version string into comparable components. + + Both '.' and '-' are treated as separators. Each component is wrapped in a + tuple: (0, int) for numeric parts, (-1, str) for non-numeric. A (0, 0) + sentinel is appended so that stable releases sort above pre-release suffixes + during Python tuple comparison (e.g., 1.2.3 >= 1.2.3.dev). Non-numeric + components are compared lexicographically (e.g., b >= a). + + Args: + version_str: Version string like "3.5.0" or "3.5.0.dev869102767". + + Returns: + Tuple of (type_order, value) pairs for comparison, where type_order is 0 + for integers and -1 for strings, followed by a (0, 0) sentinel. + """ + # Split on both '.' and '-' + parts = re.split(r"[.\-]", version_str) + return tuple( + [(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)] + ) + + +def check_version(spec: str) -> bool: + """Check if an installed package satisfies a version requirement. + + Supports operators: >=, <=, >, <, ==, != + + Version comparison rules: + - Both '.' and '-' are treated as separators + - Numeric components are compared numerically + - Non-numeric components are compared lexicographically + - Stable releases are greater than pre-releases (e.g., 1.2.3 >= 1.2.3.dev) + + Args: + spec: Version specification like "numpy>=1.20.0". + + Returns: + True if the installed version satisfies the requirement. + + Raises: + ValueError: If the spec cannot be parsed. + importlib.metadata.PackageNotFoundError: If the package is not installed. + """ + match = re.match(r"^([a-zA-Z0-9_\-]+)(>=|<=|>|<|==|!=)(.+)$", spec) + if not match: + raise ValueError( + f"Invalid version spec '{spec}'. Expected format: 'package>=version'" + ) + package_name, op, version_str = match.groups() + + required_version = _parse_version(version_str) + + try: + installed_str = importlib.metadata.version(package_name) + except importlib.metadata.PackageNotFoundError as e: + # Fallback: import the package and read __version__ + try: + import importlib as _importlib # noqa: F811 + + mod = _importlib.import_module(package_name) + installed_str = mod.__version__ + except (ImportError, AttributeError): + raise e + + installed_version = _parse_version(installed_str) + + ops = { + ">=": operator.ge, + "<=": operator.le, + ">": operator.gt, + "<": operator.lt, + "==": operator.eq, + "!=": operator.ne, + } + return ops[op](installed_version, required_version) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py index 7657346a..4ac1cb0d 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py @@ -15,6 +15,7 @@ import functools import inspect +import warnings import warp as wp @@ -144,4 +145,37 @@ def check_toolkit_driver(): wp.init() if wp.get_device().is_cuda: if not wp.is_conditional_graph_supported(): - raise RuntimeError("Minimum supported CUDA version: 12.4.") + warnings.warn( + """ + CUDA version < 12.4 detected + - graph capture may be unreliable for < 12.3 + - conditional graph nodes are not available for < 12.4 + Model.opt.graph_conditional should be set to False + """, + stacklevel=2, + ) + + +class scoped_mathdx_gemm_disabled: + """Temporarily disable Warp MathDx GEMM kernels within this scope.""" + + def __init__(self, disable: bool = True): + self._disable = disable + self._config = None + self._prev = None + + def __enter__(self): + if not self._disable: + return self + config = getattr(wp, "config", None) + if config is None or not hasattr(config, "enable_mathdx_gemm"): + return self + self._config = config + self._prev = config.enable_mathdx_gemm + self._config.enable_mathdx_gemm = False + return self + + def __exit__(self, exc_type, exc, tb): + if self._config is not None: + self._config.enable_mathdx_gemm = self._prev + return False diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml index d71512f9..37305248 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name="mujoco-warp" -version = "3.5.0" +version = "3.6.0" # TODO(team): create a distribution list authors = [ {name = "Newton Developers", email = "mujoco@deepmind.com"}, @@ -21,15 +21,16 @@ classifiers = [ "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering", ] requires-python = ">=3.10" dependencies = [ "absl-py", "etils[epath]", - "mujoco>=3.4.0", + "mujoco>=3.5.0", "numpy", - "warp-lang>=1.11.0", + "warp-lang>=1.12", ] [[tool.uv.index]] @@ -54,7 +55,7 @@ dev = [ "ruff", "pygls>=1.0.0,<2.0.0", "lsprotocol>=2023.0.1,<2024.0.0", - "mujoco>=3.4.1.dev0", + "mujoco>=3.5.0.dev0", "warp-lang>=1.11.0.dev0", ] # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 0e9287d4..2bb522f2 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -44,6 +44,9 @@ _c = mjwarp.Contact( _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) +_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( @@ -70,6 +73,7 @@ def _refit_bvh_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.flex_dim = flex_dim diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index b5822eb2..265bfba3 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -42,6 +42,10 @@ _c = mjwarp.Contact( _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) +_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( @@ -73,6 +77,7 @@ def _collision_shim( 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), @@ -135,6 +140,7 @@ def _collision_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.block_dim = block_dim @@ -163,6 +169,7 @@ def _collision_shim( _m.mesh_faceadr = mesh_faceadr _m.mesh_graph = mesh_graph _m.mesh_graphadr = mesh_graphadr + _m.mesh_octadr = mesh_octadr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr @@ -317,6 +324,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data): m.mesh_faceadr, m.mesh_graph, m.mesh_graphadr, + m.mesh_octadr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 6496f460..e3833601 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -42,6 +42,10 @@ _c = mjwarp.Contact( _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) +_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( @@ -49,7 +53,7 @@ def _forward_shim( nworld: int, M_rowadr: wp.array(dtype=int), M_rownnz: wp.array(dtype=int), - actuator_acc0: wp.array(dtype=float), + actuator_acc0: wp.array2d(dtype=float), actuator_actadr: wp.array(dtype=int), actuator_actearly: wp.array(dtype=bool), actuator_actlimited: wp.array(dtype=bool), @@ -57,7 +61,7 @@ def _forward_shim( actuator_actrange: wp.array2d(dtype=wp.vec2), actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f), actuator_biastype: wp.array(dtype=int), - actuator_cranklength: wp.array(dtype=float), + 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), @@ -67,7 +71,7 @@ def _forward_shim( 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.array(dtype=wp.vec2), + 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), @@ -94,6 +98,7 @@ def _forward_shim( 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), @@ -117,6 +122,7 @@ def _forward_shim( 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), @@ -214,6 +220,7 @@ def _forward_shim( 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), @@ -253,6 +260,7 @@ def _forward_shim( nsensortaxel: int, nsite: int, ntendon: int, + ntree: int, nu: int, nv: int, nv_pad: int, @@ -378,7 +386,7 @@ def _forward_shim( act_dot: wp.array2d(dtype=float), actuator_force: wp.array2d(dtype=float), actuator_length: wp.array2d(dtype=float), - actuator_moment: wp.array3d(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), @@ -393,7 +401,7 @@ def _forward_shim( cvel: wp.array2d(dtype=wp.spatial_vector), energy: wp.array(dtype=wp.vec2), eq_active: wp.array2d(dtype=bool), - flexedge_J: wp.array3d(dtype=float), + 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), @@ -403,11 +411,15 @@ def _forward_shim( 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), @@ -440,6 +452,7 @@ def _forward_shim( 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), @@ -466,6 +479,9 @@ def _forward_shim( 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), @@ -479,6 +495,7 @@ def _forward_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.M_rowadr = M_rowadr @@ -528,6 +545,7 @@ def _forward_shim( _m.body_rootid = body_rootid _m.body_subtreemass = body_subtreemass _m.body_tree = body_tree + _m.body_treeid = body_treeid _m.body_weldid = body_weldid _m.cam_bodyid = cam_bodyid _m.cam_fovy = cam_fovy @@ -551,6 +569,7 @@ def _forward_shim( _m.dof_parentid = dof_parentid _m.dof_solimp = dof_solimp _m.dof_solref = dof_solref + _m.dof_treeid = dof_treeid _m.dof_tri_col = dof_tri_col _m.dof_tri_row = dof_tri_row _m.eq_connect_adr = eq_connect_adr @@ -648,6 +667,7 @@ def _forward_shim( _m.mesh_normal = mesh_normal _m.mesh_normaladr = mesh_normaladr _m.mesh_normalnum = mesh_normalnum + _m.mesh_octadr = mesh_octadr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr @@ -687,6 +707,7 @@ def _forward_shim( _m.nsensortaxel = nsensortaxel _m.nsite = nsite _m.ntendon = ntendon + _m.ntree = ntree _m.nu = nu _m.nv = nv _m.nv_pad = nv_pad @@ -837,6 +858,9 @@ def _forward_shim( _d.cvel = cvel _d.efc.D = efc__D _d.efc.J = efc__J + _d.efc.J_colind = efc__J_colind + _d.efc.J_rowadr = efc__J_rowadr + _d.efc.J_rownnz = efc__J_rownnz _d.efc.Ma = efc__Ma _d.efc.aref = efc__aref _d.efc.force = efc__force @@ -859,6 +883,9 @@ def _forward_shim( _d.light_xpos = light_xpos _d.mocap_pos = mocap_pos _d.mocap_quat = mocap_quat + _d.moment_colind = moment_colind + _d.moment_rowadr = moment_rowadr + _d.moment_rownnz = moment_rownnz _d.naccdmax = naccdmax _d.nacon = nacon _d.naconmax = naconmax @@ -866,6 +893,7 @@ def _forward_shim( _d.ne = ne _d.nefc = nefc _d.nf = nf + _d.nisland = nisland _d.njmax = njmax _d.nl = nl _d.qLD = qLD @@ -899,6 +927,7 @@ def _forward_shim( _d.ten_wrapadr = ten_wrapadr _d.ten_wrapnum = ten_wrapnum _d.time = time + _d.tree_island = tree_island _d.wrap_obj = wrap_obj _d.wrap_xpos = wrap_xpos _d.xanchor = xanchor @@ -939,11 +968,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'geom_xpos': d.geom_xpos.shape, 'light_xdir': d._impl.light_xdir.shape, 'light_xpos': d._impl.light_xpos.shape, + 'moment_colind': d._impl.moment_colind.shape, + 'moment_rowadr': d._impl.moment_rowadr.shape, + 'moment_rownnz': d._impl.moment_rownnz.shape, 'nacon': d._impl.nacon.shape, 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, 'nefc': d._impl.nefc.shape, 'nf': d._impl.nf.shape, + 'nisland': d._impl.nisland.shape, 'nl': d._impl.nl.shape, 'qLD': d._impl.qLD.shape, 'qLDiagInv': d._impl.qLDiagInv.shape, @@ -972,6 +1005,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'ten_velocity': d._impl.ten_velocity.shape, 'ten_wrapadr': d._impl.ten_wrapadr.shape, 'ten_wrapnum': d._impl.ten_wrapnum.shape, + 'tree_island': d._impl.tree_island.shape, 'wrap_obj': d._impl.wrap_obj.shape, 'wrap_xpos': d._impl.wrap_xpos.shape, 'xanchor': d.xanchor.shape, @@ -997,6 +1031,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'contact__worldid': d._impl.contact__worldid.shape, 'efc__D': d._impl.efc__D.shape, 'efc__J': d._impl.efc__J.shape, + 'efc__J_colind': d._impl.efc__J_colind.shape, + 'efc__J_rowadr': d._impl.efc__J_rowadr.shape, + 'efc__J_rownnz': d._impl.efc__J_rownnz.shape, 'efc__Ma': d._impl.efc__Ma.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__force': d._impl.efc__force.shape, @@ -1010,7 +1047,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _forward_shim, - num_outputs=92, + num_outputs=100, output_dims=output_dims, vmap_method=None, in_out_argnames=set([ @@ -1038,11 +1075,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'geom_xpos', 'light_xdir', 'light_xpos', + 'moment_colind', + 'moment_rowadr', + 'moment_rownnz', 'nacon', 'ncollision', 'ne', 'nefc', 'nf', + 'nisland', 'nl', 'qLD', 'qLDiagInv', @@ -1071,6 +1112,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'ten_velocity', 'ten_wrapadr', 'ten_wrapnum', + 'tree_island', 'wrap_obj', 'wrap_xpos', 'xanchor', @@ -1096,6 +1138,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'contact__worldid', 'efc__D', 'efc__J', + 'efc__J_colind', + 'efc__J_rowadr', + 'efc__J_rownnz', 'efc__Ma', 'efc__aref', 'efc__force', @@ -1113,6 +1158,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'actuator_acc0', 'actuator_actrange', 'actuator_biasprm', + 'actuator_cranklength', 'actuator_ctrlrange', 'actuator_dynprm', 'actuator_force', @@ -1120,6 +1166,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'actuator_gainprm', 'actuator_gear', 'actuator_length', + 'actuator_lengthrange', 'body_gravcomp', 'body_inertia', 'body_invweight0', @@ -1330,6 +1377,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.body_rootid, m.body_subtreemass, m._impl.body_tree, + m.body_treeid, m.body_weldid, m.cam_bodyid, m.cam_fovy, @@ -1353,6 +1401,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.dof_parentid, m.dof_solimp, m.dof_solref, + m.dof_treeid, m._impl.dof_tri_col, m._impl.dof_tri_row, m._impl.eq_connect_adr, @@ -1450,6 +1499,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.mesh_normal, m.mesh_normaladr, m.mesh_normalnum, + m.mesh_octadr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, @@ -1489,6 +1539,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m._impl.nsensortaxel, m.nsite, m.ntendon, + m._impl.ntree, m.nu, m.nv, m._impl.nv_pad, @@ -1638,11 +1689,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.light_xpos, d.mocap_pos, d.mocap_quat, + d._impl.moment_colind, + d._impl.moment_rowadr, + d._impl.moment_rownnz, d._impl.nacon, d._impl.ncollision, d._impl.ne, d._impl.nefc, d._impl.nf, + d._impl.nisland, d._impl.nl, d._impl.qLD, d._impl.qLDiagInv, @@ -1675,6 +1730,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.ten_wrapadr, d._impl.ten_wrapnum, d.time, + d._impl.tree_island, d._impl.wrap_obj, d._impl.wrap_xpos, d.xanchor, @@ -1701,6 +1757,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.contact__worldid, d._impl.efc__D, d._impl.efc__J, + d._impl.efc__J_colind, + d._impl.efc__J_rowadr, + d._impl.efc__J_rownnz, d._impl.efc__Ma, d._impl.efc__aref, d._impl.efc__force, @@ -1737,74 +1796,82 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'geom_xpos': out[21], '_impl.light_xdir': out[22], '_impl.light_xpos': out[23], - '_impl.nacon': out[24], - '_impl.ncollision': out[25], - '_impl.ne': out[26], - '_impl.nefc': out[27], - '_impl.nf': out[28], - '_impl.nl': out[29], - '_impl.qLD': out[30], - '_impl.qLDiagInv': out[31], - '_impl.qM': out[32], - 'qacc': out[33], - 'qacc_smooth': out[34], - 'qfrc_actuator': out[35], - 'qfrc_bias': out[36], - 'qfrc_constraint': out[37], - '_impl.qfrc_damper': out[38], - 'qfrc_fluid': out[39], - 'qfrc_gravcomp': out[40], - 'qfrc_passive': out[41], - 'qfrc_smooth': out[42], - '_impl.qfrc_spring': out[43], - 'qvel': out[44], - 'sensordata': out[45], - 'site_xmat': out[46], - 'site_xpos': out[47], - '_impl.solver_niter': out[48], - '_impl.subtree_angmom': out[49], - 'subtree_com': out[50], - '_impl.subtree_linvel': out[51], - '_impl.ten_J': out[52], - 'ten_length': out[53], - '_impl.ten_velocity': out[54], - '_impl.ten_wrapadr': out[55], - '_impl.ten_wrapnum': out[56], - '_impl.wrap_obj': out[57], - '_impl.wrap_xpos': out[58], - 'xanchor': out[59], - 'xaxis': out[60], - 'ximat': out[61], - 'xipos': out[62], - 'xmat': out[63], - 'xpos': out[64], - 'xquat': out[65], - '_impl.contact__dim': out[66], - '_impl.contact__dist': out[67], - '_impl.contact__efc_address': out[68], - '_impl.contact__frame': out[69], - '_impl.contact__friction': out[70], - '_impl.contact__geom': out[71], - '_impl.contact__geomcollisionid': out[72], - '_impl.contact__includemargin': out[73], - '_impl.contact__pos': out[74], - '_impl.contact__solimp': out[75], - '_impl.contact__solref': out[76], - '_impl.contact__solreffriction': out[77], - '_impl.contact__type': out[78], - '_impl.contact__worldid': out[79], - '_impl.efc__D': out[80], - '_impl.efc__J': out[81], - '_impl.efc__Ma': out[82], - '_impl.efc__aref': out[83], - '_impl.efc__force': out[84], - '_impl.efc__frictionloss': out[85], - '_impl.efc__id': out[86], - '_impl.efc__margin': out[87], - '_impl.efc__pos': out[88], - '_impl.efc__state': out[89], - '_impl.efc__type': out[90], - '_impl.efc__vel': out[91], + '_impl.moment_colind': out[24], + '_impl.moment_rowadr': out[25], + '_impl.moment_rownnz': out[26], + '_impl.nacon': out[27], + '_impl.ncollision': out[28], + '_impl.ne': out[29], + '_impl.nefc': out[30], + '_impl.nf': out[31], + '_impl.nisland': out[32], + '_impl.nl': out[33], + '_impl.qLD': out[34], + '_impl.qLDiagInv': out[35], + '_impl.qM': out[36], + 'qacc': out[37], + 'qacc_smooth': out[38], + 'qfrc_actuator': out[39], + 'qfrc_bias': out[40], + 'qfrc_constraint': out[41], + '_impl.qfrc_damper': out[42], + 'qfrc_fluid': out[43], + 'qfrc_gravcomp': out[44], + 'qfrc_passive': out[45], + 'qfrc_smooth': out[46], + '_impl.qfrc_spring': out[47], + 'qvel': out[48], + 'sensordata': out[49], + 'site_xmat': out[50], + 'site_xpos': out[51], + '_impl.solver_niter': out[52], + '_impl.subtree_angmom': out[53], + 'subtree_com': out[54], + '_impl.subtree_linvel': out[55], + '_impl.ten_J': out[56], + 'ten_length': out[57], + '_impl.ten_velocity': out[58], + '_impl.ten_wrapadr': out[59], + '_impl.ten_wrapnum': out[60], + '_impl.tree_island': out[61], + '_impl.wrap_obj': out[62], + '_impl.wrap_xpos': out[63], + 'xanchor': out[64], + 'xaxis': out[65], + 'ximat': out[66], + 'xipos': out[67], + 'xmat': out[68], + 'xpos': out[69], + 'xquat': out[70], + '_impl.contact__dim': out[71], + '_impl.contact__dist': out[72], + '_impl.contact__efc_address': out[73], + '_impl.contact__frame': out[74], + '_impl.contact__friction': out[75], + '_impl.contact__geom': out[76], + '_impl.contact__geomcollisionid': out[77], + '_impl.contact__includemargin': out[78], + '_impl.contact__pos': out[79], + '_impl.contact__solimp': out[80], + '_impl.contact__solref': out[81], + '_impl.contact__solreffriction': out[82], + '_impl.contact__type': out[83], + '_impl.contact__worldid': out[84], + '_impl.efc__D': out[85], + '_impl.efc__J': out[86], + '_impl.efc__J_colind': out[87], + '_impl.efc__J_rowadr': out[88], + '_impl.efc__J_rownnz': out[89], + '_impl.efc__Ma': out[90], + '_impl.efc__aref': out[91], + '_impl.efc__force': out[92], + '_impl.efc__frictionloss': out[93], + '_impl.efc__id': out[94], + '_impl.efc__margin': out[95], + '_impl.efc__pos': out[96], + '_impl.efc__state': out[97], + '_impl.efc__type': out[98], + '_impl.efc__vel': out[99], }) return d @@ -1828,7 +1895,7 @@ def _step_shim( nworld: int, M_rowadr: wp.array(dtype=int), M_rownnz: wp.array(dtype=int), - actuator_acc0: wp.array(dtype=float), + actuator_acc0: wp.array2d(dtype=float), actuator_actadr: wp.array(dtype=int), actuator_actearly: wp.array(dtype=bool), actuator_actlimited: wp.array(dtype=bool), @@ -1836,7 +1903,7 @@ def _step_shim( actuator_actrange: wp.array2d(dtype=wp.vec2), actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f), actuator_biastype: wp.array(dtype=int), - actuator_cranklength: wp.array(dtype=float), + 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), @@ -1846,7 +1913,7 @@ def _step_shim( 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.array(dtype=wp.vec2), + 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), @@ -1873,6 +1940,7 @@ def _step_shim( 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), @@ -1896,6 +1964,7 @@ def _step_shim( 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), @@ -1993,6 +2062,7 @@ def _step_shim( 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), @@ -2033,6 +2103,7 @@ def _step_shim( nsensortaxel: int, nsite: int, ntendon: int, + ntree: int, nu: int, nv: int, nv_pad: int, @@ -2159,7 +2230,7 @@ def _step_shim( act_dot: wp.array2d(dtype=float), actuator_force: wp.array2d(dtype=float), actuator_length: wp.array2d(dtype=float), - actuator_moment: wp.array3d(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), @@ -2174,7 +2245,7 @@ def _step_shim( cvel: wp.array2d(dtype=wp.spatial_vector), energy: wp.array(dtype=wp.vec2), eq_active: wp.array2d(dtype=bool), - flexedge_J: wp.array3d(dtype=float), + 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), @@ -2184,11 +2255,15 @@ def _step_shim( 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), @@ -2221,6 +2296,7 @@ def _step_shim( 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), @@ -2247,6 +2323,9 @@ def _step_shim( 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), @@ -2260,6 +2339,7 @@ def _step_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.M_rowadr = M_rowadr @@ -2309,6 +2389,7 @@ def _step_shim( _m.body_rootid = body_rootid _m.body_subtreemass = body_subtreemass _m.body_tree = body_tree + _m.body_treeid = body_treeid _m.body_weldid = body_weldid _m.cam_bodyid = cam_bodyid _m.cam_fovy = cam_fovy @@ -2332,6 +2413,7 @@ def _step_shim( _m.dof_parentid = dof_parentid _m.dof_solimp = dof_solimp _m.dof_solref = dof_solref + _m.dof_treeid = dof_treeid _m.dof_tri_col = dof_tri_col _m.dof_tri_row = dof_tri_row _m.eq_connect_adr = eq_connect_adr @@ -2429,6 +2511,7 @@ def _step_shim( _m.mesh_normal = mesh_normal _m.mesh_normaladr = mesh_normaladr _m.mesh_normalnum = mesh_normalnum + _m.mesh_octadr = mesh_octadr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr @@ -2469,6 +2552,7 @@ def _step_shim( _m.nsensortaxel = nsensortaxel _m.nsite = nsite _m.ntendon = ntendon + _m.ntree = ntree _m.nu = nu _m.nv = nv _m.nv_pad = nv_pad @@ -2620,6 +2704,9 @@ def _step_shim( _d.cvel = cvel _d.efc.D = efc__D _d.efc.J = efc__J + _d.efc.J_colind = efc__J_colind + _d.efc.J_rowadr = efc__J_rowadr + _d.efc.J_rownnz = efc__J_rownnz _d.efc.Ma = efc__Ma _d.efc.aref = efc__aref _d.efc.force = efc__force @@ -2642,6 +2729,9 @@ def _step_shim( _d.light_xpos = light_xpos _d.mocap_pos = mocap_pos _d.mocap_quat = mocap_quat + _d.moment_colind = moment_colind + _d.moment_rowadr = moment_rowadr + _d.moment_rownnz = moment_rownnz _d.naccdmax = naccdmax _d.nacon = nacon _d.naconmax = naconmax @@ -2649,6 +2739,7 @@ def _step_shim( _d.ne = ne _d.nefc = nefc _d.nf = nf + _d.nisland = nisland _d.njmax = njmax _d.nl = nl _d.qLD = qLD @@ -2682,6 +2773,7 @@ def _step_shim( _d.ten_wrapadr = ten_wrapadr _d.ten_wrapnum = ten_wrapnum _d.time = time + _d.tree_island = tree_island _d.wrap_obj = wrap_obj _d.wrap_xpos = wrap_xpos _d.xanchor = xanchor @@ -2723,11 +2815,15 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'geom_xpos': d.geom_xpos.shape, 'light_xdir': d._impl.light_xdir.shape, 'light_xpos': d._impl.light_xpos.shape, + 'moment_colind': d._impl.moment_colind.shape, + 'moment_rowadr': d._impl.moment_rowadr.shape, + 'moment_rownnz': d._impl.moment_rownnz.shape, 'nacon': d._impl.nacon.shape, 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, 'nefc': d._impl.nefc.shape, 'nf': d._impl.nf.shape, + 'nisland': d._impl.nisland.shape, 'nl': d._impl.nl.shape, 'qLD': d._impl.qLD.shape, 'qLDiagInv': d._impl.qLDiagInv.shape, @@ -2759,6 +2855,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'ten_wrapadr': d._impl.ten_wrapadr.shape, 'ten_wrapnum': d._impl.ten_wrapnum.shape, 'time': d.time.shape, + 'tree_island': d._impl.tree_island.shape, 'wrap_obj': d._impl.wrap_obj.shape, 'wrap_xpos': d._impl.wrap_xpos.shape, 'xanchor': d.xanchor.shape, @@ -2784,6 +2881,9 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'contact__worldid': d._impl.contact__worldid.shape, 'efc__D': d._impl.efc__D.shape, 'efc__J': d._impl.efc__J.shape, + 'efc__J_colind': d._impl.efc__J_colind.shape, + 'efc__J_rowadr': d._impl.efc__J_rowadr.shape, + 'efc__J_rownnz': d._impl.efc__J_rownnz.shape, 'efc__Ma': d._impl.efc__Ma.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__force': d._impl.efc__force.shape, @@ -2797,7 +2897,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _step_shim, - num_outputs=96, + num_outputs=104, output_dims=output_dims, vmap_method=None, in_out_argnames=set([ @@ -2826,11 +2926,15 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'geom_xpos', 'light_xdir', 'light_xpos', + 'moment_colind', + 'moment_rowadr', + 'moment_rownnz', 'nacon', 'ncollision', 'ne', 'nefc', 'nf', + 'nisland', 'nl', 'qLD', 'qLDiagInv', @@ -2862,6 +2966,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'ten_wrapadr', 'ten_wrapnum', 'time', + 'tree_island', 'wrap_obj', 'wrap_xpos', 'xanchor', @@ -2887,6 +2992,9 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'contact__worldid', 'efc__D', 'efc__J', + 'efc__J_colind', + 'efc__J_rowadr', + 'efc__J_rownnz', 'efc__Ma', 'efc__aref', 'efc__force', @@ -2904,6 +3012,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'actuator_acc0', 'actuator_actrange', 'actuator_biasprm', + 'actuator_cranklength', 'actuator_ctrlrange', 'actuator_dynprm', 'actuator_force', @@ -2911,6 +3020,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'actuator_gainprm', 'actuator_gear', 'actuator_length', + 'actuator_lengthrange', 'body_gravcomp', 'body_inertia', 'body_invweight0', @@ -3125,6 +3235,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.body_rootid, m.body_subtreemass, m._impl.body_tree, + m.body_treeid, m.body_weldid, m.cam_bodyid, m.cam_fovy, @@ -3148,6 +3259,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.dof_parentid, m.dof_solimp, m.dof_solref, + m.dof_treeid, m._impl.dof_tri_col, m._impl.dof_tri_row, m._impl.eq_connect_adr, @@ -3245,6 +3357,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.mesh_normal, m.mesh_normaladr, m.mesh_normalnum, + m.mesh_octadr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, @@ -3285,6 +3398,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m._impl.nsensortaxel, m.nsite, m.ntendon, + m._impl.ntree, m.nu, m.nv, m._impl.nv_pad, @@ -3435,11 +3549,15 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.light_xpos, d.mocap_pos, d.mocap_quat, + d._impl.moment_colind, + d._impl.moment_rowadr, + d._impl.moment_rownnz, d._impl.nacon, d._impl.ncollision, d._impl.ne, d._impl.nefc, d._impl.nf, + d._impl.nisland, d._impl.nl, d._impl.qLD, d._impl.qLDiagInv, @@ -3472,6 +3590,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.ten_wrapadr, d._impl.ten_wrapnum, d.time, + d._impl.tree_island, d._impl.wrap_obj, d._impl.wrap_xpos, d.xanchor, @@ -3498,6 +3617,9 @@ def _step_jax_impl(m: types.Model, d: types.Data): d._impl.contact__worldid, d._impl.efc__D, d._impl.efc__J, + d._impl.efc__J_colind, + d._impl.efc__J_rowadr, + d._impl.efc__J_rownnz, d._impl.efc__Ma, d._impl.efc__aref, d._impl.efc__force, @@ -3535,77 +3657,85 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'geom_xpos': out[22], '_impl.light_xdir': out[23], '_impl.light_xpos': out[24], - '_impl.nacon': out[25], - '_impl.ncollision': out[26], - '_impl.ne': out[27], - '_impl.nefc': out[28], - '_impl.nf': out[29], - '_impl.nl': out[30], - '_impl.qLD': out[31], - '_impl.qLDiagInv': out[32], - '_impl.qM': out[33], - 'qacc': out[34], - 'qacc_smooth': out[35], - 'qacc_warmstart': out[36], - 'qfrc_actuator': out[37], - 'qfrc_bias': out[38], - 'qfrc_constraint': out[39], - '_impl.qfrc_damper': out[40], - 'qfrc_fluid': out[41], - 'qfrc_gravcomp': out[42], - 'qfrc_passive': out[43], - 'qfrc_smooth': out[44], - '_impl.qfrc_spring': out[45], - 'qpos': out[46], - 'qvel': out[47], - 'sensordata': out[48], - 'site_xmat': out[49], - 'site_xpos': out[50], - '_impl.solver_niter': out[51], - '_impl.subtree_angmom': out[52], - 'subtree_com': out[53], - '_impl.subtree_linvel': out[54], - '_impl.ten_J': out[55], - 'ten_length': out[56], - '_impl.ten_velocity': out[57], - '_impl.ten_wrapadr': out[58], - '_impl.ten_wrapnum': out[59], - 'time': out[60], - '_impl.wrap_obj': out[61], - '_impl.wrap_xpos': out[62], - 'xanchor': out[63], - 'xaxis': out[64], - 'ximat': out[65], - 'xipos': out[66], - 'xmat': out[67], - 'xpos': out[68], - 'xquat': out[69], - '_impl.contact__dim': out[70], - '_impl.contact__dist': out[71], - '_impl.contact__efc_address': out[72], - '_impl.contact__frame': out[73], - '_impl.contact__friction': out[74], - '_impl.contact__geom': out[75], - '_impl.contact__geomcollisionid': out[76], - '_impl.contact__includemargin': out[77], - '_impl.contact__pos': out[78], - '_impl.contact__solimp': out[79], - '_impl.contact__solref': out[80], - '_impl.contact__solreffriction': out[81], - '_impl.contact__type': out[82], - '_impl.contact__worldid': out[83], - '_impl.efc__D': out[84], - '_impl.efc__J': out[85], - '_impl.efc__Ma': out[86], - '_impl.efc__aref': out[87], - '_impl.efc__force': out[88], - '_impl.efc__frictionloss': out[89], - '_impl.efc__id': out[90], - '_impl.efc__margin': out[91], - '_impl.efc__pos': out[92], - '_impl.efc__state': out[93], - '_impl.efc__type': out[94], - '_impl.efc__vel': out[95], + '_impl.moment_colind': out[25], + '_impl.moment_rowadr': out[26], + '_impl.moment_rownnz': out[27], + '_impl.nacon': out[28], + '_impl.ncollision': out[29], + '_impl.ne': out[30], + '_impl.nefc': out[31], + '_impl.nf': out[32], + '_impl.nisland': out[33], + '_impl.nl': out[34], + '_impl.qLD': out[35], + '_impl.qLDiagInv': out[36], + '_impl.qM': out[37], + 'qacc': out[38], + 'qacc_smooth': out[39], + 'qacc_warmstart': out[40], + 'qfrc_actuator': out[41], + 'qfrc_bias': out[42], + 'qfrc_constraint': out[43], + '_impl.qfrc_damper': out[44], + 'qfrc_fluid': out[45], + 'qfrc_gravcomp': out[46], + 'qfrc_passive': out[47], + 'qfrc_smooth': out[48], + '_impl.qfrc_spring': out[49], + 'qpos': out[50], + 'qvel': out[51], + 'sensordata': out[52], + 'site_xmat': out[53], + 'site_xpos': out[54], + '_impl.solver_niter': out[55], + '_impl.subtree_angmom': out[56], + 'subtree_com': out[57], + '_impl.subtree_linvel': out[58], + '_impl.ten_J': out[59], + 'ten_length': out[60], + '_impl.ten_velocity': out[61], + '_impl.ten_wrapadr': out[62], + '_impl.ten_wrapnum': out[63], + 'time': out[64], + '_impl.tree_island': out[65], + '_impl.wrap_obj': out[66], + '_impl.wrap_xpos': out[67], + 'xanchor': out[68], + 'xaxis': out[69], + 'ximat': out[70], + 'xipos': out[71], + 'xmat': out[72], + 'xpos': out[73], + 'xquat': out[74], + '_impl.contact__dim': out[75], + '_impl.contact__dist': out[76], + '_impl.contact__efc_address': out[77], + '_impl.contact__frame': out[78], + '_impl.contact__friction': out[79], + '_impl.contact__geom': out[80], + '_impl.contact__geomcollisionid': out[81], + '_impl.contact__includemargin': out[82], + '_impl.contact__pos': out[83], + '_impl.contact__solimp': out[84], + '_impl.contact__solref': out[85], + '_impl.contact__solreffriction': out[86], + '_impl.contact__type': out[87], + '_impl.contact__worldid': out[88], + '_impl.efc__D': out[89], + '_impl.efc__J': out[90], + '_impl.efc__J_colind': out[91], + '_impl.efc__J_rowadr': out[92], + '_impl.efc__J_rownnz': out[93], + '_impl.efc__Ma': out[94], + '_impl.efc__aref': out[95], + '_impl.efc__force': out[96], + '_impl.efc__frictionloss': out[97], + '_impl.efc__id': out[98], + '_impl.efc__margin': out[99], + '_impl.efc__pos': out[100], + '_impl.efc__state': out[101], + '_impl.efc__type': out[102], + '_impl.efc__vel': out[103], }) return d diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 3eb936a7..15e874e3 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -181,7 +181,15 @@ class ForwardTest(parameterized.TestCase): d.moment_rowadr, d.moment_colind, ) - tu.assert_eq(dx._impl.actuator_moment, actuator_moment, 'actuator_moment') + warp_actuator_moment = np.zeros((m.nu, m.nv)) + mujoco.mju_sparse2dense( + warp_actuator_moment, + np.asarray(dx._impl.actuator_moment), + np.asarray(dx._impl.moment_rownnz), + np.asarray(dx._impl.moment_rowadr), + np.asarray(dx._impl.moment_colind), + ) + tu.assert_eq(warp_actuator_moment, actuator_moment, 'actuator_moment') # fwd_velocity tu.assert_attr_eq(dx._impl, d, 'actuator_velocity') diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 446431e7..0bab51a3 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -44,6 +44,9 @@ _c = mjwarp.Contact( _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) +_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( @@ -81,6 +84,7 @@ def _render_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.cam_fovy = cam_fovy diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 310b3562..eeef2d92 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -42,6 +42,9 @@ _c = mjwarp.Contact( _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) +_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( @@ -92,6 +95,7 @@ def _kinematics_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.body_branch_start = body_branch_start @@ -332,6 +336,7 @@ def _tendon_shim( ): _m.stat = _s _m.opt = _o + _m.callback = _cb _d.efc = _e _d.contact = _c _m.body_parentid = body_parentid diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 56ee25f8..ff85189f 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -25,11 +25,20 @@ from mujoco.mjx._src import dataclasses as mjx_dataclasses import numpy as np if typing.TYPE_CHECKING: GraphMode = int + + @dataclasses.dataclass + class Callback: + pass + else: try: from warp._src.jax_experimental.ffi import GraphMode + from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types + + Callback = mjwp_types.Callback except ImportError: GraphMode = int + Callback = None PyTreeNode = mjx_dataclasses.PyTreeNode @dataclasses.dataclass(frozen=True) @@ -43,7 +52,6 @@ class TileSet: adr: address of each tile in the set size: size of all the tiles in this set """ - adr: np.ndarray size: int @@ -122,6 +130,7 @@ class ModelWarp(PyTreeNode): body_branches: np.ndarray body_fluid_ellipsoid: np.ndarray body_tree: Tuple[np.ndarray, ...] + callback: Callback cam_projection: np.ndarray collision_sensor_adr: np.ndarray dof_tri_col: np.ndarray @@ -174,6 +183,7 @@ class ModelWarp(PyTreeNode): mesh_polyvertadr: np.ndarray mesh_polyvertnum: np.ndarray mocap_bodyid: np.ndarray + nJfe: int nacttrnbody: int nbranch: int nflex: int @@ -270,6 +280,9 @@ class DataWarp(PyTreeNode): crb: jax.Array efc__D: jax.Array efc__J: jax.Array + efc__J_colind: jax.Array + efc__J_rowadr: jax.Array + efc__J_rownnz: jax.Array efc__Ma: jax.Array efc__aref: jax.Array efc__force: jax.Array @@ -287,6 +300,9 @@ class DataWarp(PyTreeNode): flexvert_xpos: jax.Array light_xdir: jax.Array light_xpos: jax.Array + moment_colind: jax.Array + moment_rowadr: jax.Array + moment_rownnz: jax.Array naccdmax: int nacon: jax.Array naconmax: int @@ -294,7 +310,9 @@ class DataWarp(PyTreeNode): ne: jax.Array nefc: jax.Array nf: jax.Array + nisland: jax.Array njmax: int + njmax_pad: int nl: jax.Array nworld: int qLD: jax.Array @@ -309,6 +327,7 @@ class DataWarp(PyTreeNode): ten_velocity: jax.Array ten_wrapadr: jax.Array ten_wrapnum: jax.Array + tree_island: jax.Array wrap_obj: jax.Array wrap_xpos: jax.Array shape = property(lambda self: self.cacc.shape) @@ -332,6 +351,7 @@ DATA_NON_VMAP = { 'naconmax', 'ncollision', 'njmax', + 'njmax_pad', 'nworld', } @@ -365,7 +385,7 @@ _NDIM = { 'act_dot': 2, 'actuator_force': 2, 'actuator_length': 2, - 'actuator_moment': 3, + 'actuator_moment': 2, 'actuator_velocity': 2, 'cacc': 3, 'cam_xmat': 4, @@ -394,6 +414,9 @@ _NDIM = { 'cvel': 3, 'efc__D': 2, 'efc__J': 3, + 'efc__J_colind': 3, + 'efc__J_rowadr': 2, + 'efc__J_rownnz': 2, 'efc__Ma': 2, 'efc__aref': 2, 'efc__force': 2, @@ -406,7 +429,7 @@ _NDIM = { 'efc__vel': 2, 'energy': 2, 'eq_active': 2, - 'flexedge_J': 3, + 'flexedge_J': 2, 'flexedge_length': 2, 'flexedge_velocity': 2, 'flexvert_xpos': 3, @@ -416,6 +439,9 @@ _NDIM = { 'light_xpos': 3, 'mocap_pos': 3, 'mocap_quat': 3, + 'moment_colind': 2, + 'moment_rowadr': 2, + 'moment_rownnz': 2, 'naccdmax': 0, 'nacon': 1, 'naconmax': 0, @@ -423,7 +449,9 @@ _NDIM = { 'ne': 1, 'nefc': 1, 'nf': 1, + 'nisland': 1, 'njmax': 0, + 'njmax_pad': 0, 'nl': 1, 'nworld': 0, 'qLD': 3, @@ -458,6 +486,7 @@ _NDIM = { 'ten_wrapadr': 2, 'ten_wrapnum': 2, 'time': 1, + 'tree_island': 2, 'wrap_obj': 3, 'wrap_xpos': 3, 'xanchor': 3, @@ -473,7 +502,7 @@ _NDIM = { 'M_colind': 1, 'M_rowadr': 1, 'M_rownnz': 1, - 'actuator_acc0': 1, + 'actuator_acc0': 2, 'actuator_actadr': 1, 'actuator_actearly': 1, 'actuator_actlimited': 1, @@ -481,7 +510,7 @@ _NDIM = { 'actuator_actrange': 3, 'actuator_biasprm': 3, 'actuator_biastype': 1, - 'actuator_cranklength': 1, + 'actuator_cranklength': 2, 'actuator_ctrllimited': 1, 'actuator_ctrlrange': 3, 'actuator_dynprm': 3, @@ -491,7 +520,7 @@ _NDIM = { 'actuator_gainprm': 3, 'actuator_gaintype': 1, 'actuator_gear': 3, - 'actuator_lengthrange': 2, + 'actuator_lengthrange': 3, 'actuator_trnid': 2, 'actuator_trntype': 1, 'actuator_trntype_body_adr': 1, @@ -670,6 +699,7 @@ _NDIM = { 'mesh_normal': 2, 'mesh_normaladr': 1, 'mesh_normalnum': 1, + 'mesh_octadr': 1, 'mesh_polyadr': 1, 'mesh_polymap': 1, 'mesh_polymapadr': 1, @@ -685,6 +715,8 @@ _NDIM = { 'mesh_vertnum': 1, 'mocap_bodyid': 1, 'nC': 0, + 'nJfe': 0, + 'nJmom': 0, 'nM': 0, 'na': 0, 'nacttrnbody': 0, @@ -924,6 +956,9 @@ _BATCH_DIM = { 'cvel': True, 'efc__D': True, 'efc__J': True, + 'efc__J_colind': True, + 'efc__J_rowadr': True, + 'efc__J_rownnz': True, 'efc__Ma': True, 'efc__aref': True, 'efc__force': True, @@ -946,6 +981,9 @@ _BATCH_DIM = { 'light_xpos': True, 'mocap_pos': True, 'mocap_quat': True, + 'moment_colind': True, + 'moment_rowadr': True, + 'moment_rownnz': True, 'naccdmax': False, 'nacon': False, 'naconmax': False, @@ -953,7 +991,9 @@ _BATCH_DIM = { 'ne': True, 'nefc': True, 'nf': True, + 'nisland': True, 'njmax': False, + 'njmax_pad': False, 'nl': True, 'nworld': False, 'qLD': True, @@ -988,6 +1028,7 @@ _BATCH_DIM = { 'ten_wrapadr': True, 'ten_wrapnum': True, 'time': True, + 'tree_island': True, 'wrap_obj': True, 'wrap_xpos': True, 'xanchor': True, @@ -1003,7 +1044,7 @@ _BATCH_DIM = { 'M_colind': False, 'M_rowadr': False, 'M_rownnz': False, - 'actuator_acc0': False, + 'actuator_acc0': True, 'actuator_actadr': False, 'actuator_actearly': False, 'actuator_actlimited': False, @@ -1011,7 +1052,7 @@ _BATCH_DIM = { 'actuator_actrange': True, 'actuator_biasprm': True, 'actuator_biastype': False, - 'actuator_cranklength': False, + 'actuator_cranklength': True, 'actuator_ctrllimited': False, 'actuator_ctrlrange': True, 'actuator_dynprm': True, @@ -1021,7 +1062,7 @@ _BATCH_DIM = { 'actuator_gainprm': True, 'actuator_gaintype': False, 'actuator_gear': True, - 'actuator_lengthrange': False, + 'actuator_lengthrange': True, 'actuator_trnid': False, 'actuator_trntype': False, 'actuator_trntype_body_adr': False, @@ -1200,6 +1241,7 @@ _BATCH_DIM = { 'mesh_normal': False, 'mesh_normaladr': False, 'mesh_normalnum': False, + 'mesh_octadr': False, 'mesh_polyadr': False, 'mesh_polymap': False, 'mesh_polymapadr': False, @@ -1215,6 +1257,8 @@ _BATCH_DIM = { 'mesh_vertnum': False, 'mocap_bodyid': False, 'nC': False, + 'nJfe': False, + 'nJmom': False, 'nM': False, 'na': False, 'nacttrnbody': False,