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

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 885715192
Change-Id: I13fcb8312c1515bb1bd2d6cdfd0d256e37cb74de
This commit is contained in:
Taylor Howell
2026-03-18 11:57:39 -07:00
committed by Copybara-Service
parent eec6e43a73
commit 57bdf89d07
33 changed files with 6038 additions and 3457 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+3
View File
@@ -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
mjx/mujoco/mjx/_src/types.py
+2 -2
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
Vendored
+2
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
Vendored
+2 -2
View File
@@ -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])
mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
Vendored
+2 -1
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+20 -14
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+372
View File
@@ -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),
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+29 -39
View File
@@ -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)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+16 -14
View File
@@ -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:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+8 -303
View File
@@ -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:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+217 -174
View File
@@ -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,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+1891 -1401
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+51 -66
View File
@@ -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),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+226 -135
View File
@@ -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):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+502 -186
View File
@@ -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):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+89 -3
View File
@@ -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],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+48 -45
View File
@@ -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)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+3 -3
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+7 -1
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+239 -133
View File
@@ -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(
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+533 -287
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+1081 -398
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+194 -84
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
Vendored
+97
View File
@@ -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)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
Vendored
+35 -1
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+5 -4
View File
@@ -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
mjx/mujoco/mjx/warp/bvh.py
+4
View File
@@ -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
mjx/mujoco/mjx/warp/collision_driver.py
+8
View File
@@ -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,
mjx/mujoco/mjx/warp/forward.py
+281 -151
View File
@@ -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
mjx/mujoco/mjx/warp/forward_test.py
+9 -1
View File
@@ -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')
mjx/mujoco/mjx/warp/render.py
+4
View File
@@ -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
mjx/mujoco/mjx/warp/smooth.py
+5
View File
@@ -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
mjx/mujoco/mjx/warp/types.py
+53 -9
View File
@@ -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,