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: 923922783
Change-Id: I0b407347a70a0e2024aae7ed5bab1592dc40a074
This commit is contained in:
Taylor Howell
2026-05-30 07:43:56 -07:00
committed by Copybara-Service
parent b9c7a4b81a
commit 50e823e91c
31 changed files with 8980 additions and 2696 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+11 -2
View File
@@ -272,6 +272,10 @@ def _put_option(
if impl == types.Impl.WARP:
if not mjxw.mjwp_io.ENABLE_ISLANDS:
fields['disableflags'] = types.DisableBit(
fields['disableflags'] | mjwp_types.DisableBit.ISLAND
)
impl_fields = {k: _wp_to_np_type(v) for k, v in impl_fields.items()}
return types.Option(**fields, _impl=mjxw.types.OptionWarp(**impl_fields))
@@ -1272,8 +1276,11 @@ def _get_data_into_warp(
else:
value = getattr(d_i, field.name)
if field.name in ('ne', 'nl', 'nf'):
pass
if field.name in ('ne', 'nl', 'nf', 'nisland', 'nidof'):
if isinstance(value, np.ndarray) and value.size == 0:
value = 0
else:
value = int(value)
elif field.name in ('nefc', 'ncon'):
value = {'nefc': nefc, 'ncon': ncon}[field.name]
elif field.name.endswith('xmat') or field.name == 'ximat':
@@ -1287,9 +1294,11 @@ def _get_data_into_warp(
'contact',
'qM',
'qLD',
'qLU',
'qLDiagInv',
'ten_J',
'flexedge_J',
'M',
):
continue
if field.name.startswith('efc_'):
mjx/mujoco/mjx/_src/io_test.py
+1 -1
View File
@@ -530,7 +530,7 @@ class DataIOTest(parameterized.TestCase):
elif impl == 'warp':
qm = np.zeros((m.nv, m.nv), dtype=np.float64)
mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
np.testing.assert_allclose(dx._impl.qM, qm)
np.testing.assert_allclose(dx._impl.M[:m.nv, :m.nv], qm)
# TODO(taylorhowell): test efc__J
np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
Vendored
+4
View File
@@ -47,6 +47,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit
from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as rungekutta4
from mujoco.mjx.third_party.mujoco_warp._src.forward import step1 as step1
from mujoco.mjx.third_party.mujoco_warp._src.forward import step2 as step2
from mujoco.mjx.third_party.mujoco_warp._src.history import init_ctrl_history as init_ctrl_history
from mujoco.mjx.third_party.mujoco_warp._src.history import init_sensor_history as init_sensor_history
from mujoco.mjx.third_party.mujoco_warp._src.history import read_ctrl as read_ctrl
from mujoco.mjx.third_party.mujoco_warp._src.history import read_sensor as read_sensor
from mujoco.mjx.third_party.mujoco_warp._src.inverse import inverse as inverse
from mujoco.mjx.third_party.mujoco_warp._src.io import create_render_context as create_render_context
from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into as get_data_into
mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
Vendored
+73 -36
View File
@@ -19,48 +19,77 @@ import warp as wp
@lru_cache(maxsize=None)
def create_blocked_cholesky_func(block_size: int):
def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
@wp.func
def blocked_cholesky_func(
def blocked_cholesky_factorize_solve_func(
# In:
A: wp.array2d[float],
b: wp.array2d[float],
matrix_size: int,
# Out:
L: wp.array2d[float],
U: wp.array2d[float],
x: wp.array2d[float],
):
"""Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
"""Block Cholesky factorization and solve while keeping the forward RHS live."""
rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
It returns a lower-triangular matrix L such that A = L L^T.
"""
# Process the matrix in blocks along its leading dimension.
for k in range(0, matrix_size, block_size):
end = k + block_size
rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(k, 0))
# Load current diagonal block A[k:end, k:end]
# and update with contributions from previously computed blocks.
A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared")
A_kk_tile = wp.tile_load(
A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
)
for j in range(0, k, block_size):
L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
U_block = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
)
wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
# Compute the Cholesky factorization for the block
wp.tile_cholesky_inplace(A_kk_tile)
wp.tile_store(L, A_kk_tile, offset=(k, k))
y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
wp.tile_store(U, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), rhs_view)
# Process the blocks below the current block
for i in range(end, matrix_size, block_size):
A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared")
A_ki_tile = wp.tile_load(
A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
)
for j in range(0, k, block_size):
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
U_jk_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
)
U_ji_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
wp.tile_lower_solve_inplace(A_kk_tile, wp.tile_transpose(A_ik_tile))
wp.tile_store(L, A_ik_tile, offset=(i, k))
wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
wp.tile_store(U, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
return blocked_cholesky_func
for i in range(matrix_size - block_size, -1, -block_size):
i_end = i + block_size
tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(i_end, matrix_size, block_size):
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
)
x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_upper_solve_inplace(U_tile, tmp_tile)
wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
return blocked_cholesky_factorize_solve_func
@lru_cache(maxsize=None)
@@ -68,41 +97,49 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
@wp.func
def blocked_cholesky_solve_func(
# In:
L: wp.array2d[float],
U: wp.array2d[float],
b: wp.array2d[float],
matrix_size: int,
# Out:
x: wp.array2d[float],
):
"""Block Cholesky factorization and solve.
"""Block Cholesky solve.
Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
Solves A x = b given the Cholesky factor U (A = U^T U) using blocked forward and backward
substitution.
"""
rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
# Forward substitution: solve L y = b
# Forward substitution: solve U^T y = b
for i in range(0, matrix_size, block_size):
rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(0, i, block_size):
L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
U_block = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
)
y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
wp.tile_lower_solve_inplace(L_tile, rhs_view)
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_lower_solve_inplace(wp.tile_transpose(U_tile), rhs_view)
# Backward substitution: solve L^T x = y
# Backward substitution: solve U x = y
for i in range(matrix_size - block_size, -1, -block_size):
i_end = i + block_size
tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(i_end, matrix_size, block_size):
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
)
x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
wp.tile_upper_solve_inplace(U_tile, tmp_tile)
wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
Vendored
+23 -6
View File
@@ -25,6 +25,7 @@ from absl import flags
from etils import epath
import mujoco.mjx.third_party.mujoco_warp as mjw
from mujoco.mjx.third_party.mujoco_warp._src import io
from mujoco.mjx.third_party.mujoco_warp._src import warp_util
from mujoco.mjx.third_party.mujoco_warp._src.io import load_trajectory
from mujoco.mjx.third_party.mujoco_warp._src.io import override_model
@@ -42,6 +43,12 @@ KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulatio
EVENT_TRACE = flags.DEFINE_bool("event_trace", False, "print an event trace report")
NOISE_STD = flags.DEFINE_float("noise_std", 0.01, "add noise to ctrl signal (standard deviation)")
NOISE_RATE = flags.DEFINE_float("noise_rate", 0.1, "add noise to ctrl signal (noise rate)")
ENABLE_ISLANDS = flags.DEFINE_bool(
"enable_islands",
False,
"Enable constraint islands solver",
)
DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
REPLAY = flags.DEFINE_string("replay", None, "NPZ file with ctrl sequence to replay")
@@ -52,6 +59,12 @@ RENDER_RGB = flags.DEFINE_bool("render_rgb", True, "render RGB image")
RENDER_DEPTH = flags.DEFINE_bool("render_depth", True, "render depth image")
RENDER_TEXTURES = flags.DEFINE_bool("render_textures", True, "use textures")
RENDER_SHADOWS = flags.DEFINE_bool("render_shadows", False, "use shadows")
RENDER_BACKFACE_CULLING = flags.DEFINE_bool(
"render_backface_culling",
True,
"enable renderer backface culling (RenderContext.enable_backface_culling)",
)
RENDER_SKYBOX = flags.DEFINE_bool("render_skybox", True, "render skybox texture if available")
def load_model(path: epath.Path) -> mujoco.MjModel:
@@ -128,6 +141,8 @@ def init_structs(
fn: Callable[..., None], mjm: mujoco.MjModel
) -> Tuple[mjw.Model, mjw.Data, mjw.RenderContext | None, list[np.ndarray] | None]:
"""Initialize device structs."""
io.ENABLE_ISLANDS = ENABLE_ISLANDS.value
mjd = mujoco.MjData(mjm)
ctrls = None
if REPLAY.value:
@@ -152,12 +167,14 @@ def init_structs(
rc = mjw.create_render_context(
mjm,
NWORLD.value,
(RENDER_WIDTH.value, RENDER_HEIGHT.value),
RENDER_RGB.value,
RENDER_DEPTH.value,
RENDER_TEXTURES.value,
RENDER_SHADOWS.value,
nworld=NWORLD.value,
cam_res=(RENDER_WIDTH.value, RENDER_HEIGHT.value),
render_rgb=RENDER_RGB.value,
render_depth=RENDER_DEPTH.value,
use_textures=RENDER_TEXTURES.value,
use_shadows=RENDER_SHADOWS.value,
enable_backface_culling=RENDER_BACKFACE_CULLING.value,
render_skybox=RENDER_SKYBOX.value,
)
return m, d, rc, ctrls
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+34 -18
View File
@@ -23,6 +23,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_param
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 epa_phase
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import gjk_phase
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
@@ -35,7 +37,7 @@ 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 _NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
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
@@ -715,6 +717,7 @@ def ccd_kernel_builder(
opt_ccd_tolerance: wp.array[float],
# Data in:
naconmax_in: int,
naccdmax_in: int,
# In:
epa_vert_in: wp.array2d[wp.vec3],
epa_vert_index_in: wp.array2d[int],
@@ -737,7 +740,7 @@ def ccd_kernel_builder(
geom2: Geom,
geoms: wp.vec2i,
worldid: int,
ccdid: int,
nccd_in: wp.array[int],
margin: float,
gap: float,
condim: int,
@@ -774,30 +777,47 @@ def ccd_kernel_builder(
if is_collision_sensor:
cutoff = 1.0e32
else:
if wp.static(_NEW_GAP_SEMANTICS):
if wp.static(NEW_GAP_SEMANTICS):
cutoff = gap
else:
cutoff = 0.0
dist, ncollision, w1, w2, multiccd_idx = ccd(
needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
cutoff,
gjk_iterations,
epa_iterations,
geom1,
geom2,
geomtype1,
geomtype2,
x1,
x2,
epa_vert_in[ccdid],
epa_vert_index_in[ccdid],
epa_face_in[ccdid],
epa_pr_in[ccdid],
epa_norm2_in[ccdid],
epa_horizon_in[ccdid],
)
if wp.static(_NEW_GAP_SEMANTICS):
ccdid = int(-1)
multiccd_idx = int(-1)
if needs_epa:
ccdid = wp.atomic_add(nccd_in, geomgeomid, 1)
if ccdid >= naccdmax_in:
wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
return 0
dist, ncollision, w1, w2, multiccd_idx = epa_phase(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
epa_iterations,
gjk_result,
geom1,
geom2,
geomtype1,
geomtype2,
epa_vert_in[ccdid],
epa_vert_index_in[ccdid],
epa_face_in[ccdid],
epa_pr_in[ccdid],
epa_norm2_in[ccdid],
epa_horizon_in[ccdid],
)
if wp.static(NEW_GAP_SEMANTICS):
if dist >= gap and pairid[1] == -1:
return 0
else:
@@ -990,11 +1010,6 @@ def ccd_kernel_builder(
if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
return
ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
if ccdid >= naccdmax_in:
wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
return
worldid = collision_worldid_in[collisionid]
_, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
@@ -1046,6 +1061,7 @@ def ccd_kernel_builder(
eval_ccd_write_contact(
opt_ccd_tolerance,
naconmax_in,
naccdmax_in,
epa_vert_in,
epa_vert_index_in,
epa_face_in,
@@ -1067,7 +1083,7 @@ def ccd_kernel_builder(
geom2,
geoms,
worldid,
ccdid,
nccd_in,
margin,
gap,
condim,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+2 -2
View File
@@ -22,8 +22,8 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
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
@@ -219,7 +219,7 @@ def write_contact(
contact_frame_out[cid] = frame_in
contact_geom_out[cid] = geoms_in
contact_worldid_out[cid] = worldid_in
if wp.static(_NEW_GAP_SEMANTICS):
if wp.static(NEW_GAP_SEMANTICS):
includemargin = margin_in
else:
includemargin = margin_in - gap_in
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+8 -2
View File
@@ -375,6 +375,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
return upper
@cache_kernel
def _sap_project(opt_broadphase: int):
@wp.kernel(module="unique", enable_backward=False)
def sap_project(
@@ -531,6 +532,7 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
return kernel
@cache_kernel
def _segmented_sort(tile_size: int):
@wp.kernel(module="unique")
def segmented_sort(
@@ -629,7 +631,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
# assumes each geom has 5 other geoms (batched over all worlds)
nsweep = 5 * nworldgeom
wp.launch(
kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
kernel=_sap_broadphase(
m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
),
dim=nsweep,
inputs=[
m.ngeom,
@@ -718,7 +722,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
"""
wp.launch(
_nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
_nxn_broadphase(
m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[
m.geom_type,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+80 -17
View File
@@ -2135,8 +2135,12 @@ def multicontact(
# face1 is an edge; clip face1 against face2
if is_edge_contact_geom1:
approx_dir = wp.norm_l2(dir) * n2[j]
return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
approx_dir = -wp.norm_l2(dir) * n2[j]
nclipped, clipped1, clipped2 = _polygon_clip(
plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped
)
# the faces were flipped in calling _polygon_clip so we need to flip them back
return nclipped, clipped2, clipped1
# face2 is an edge; clip face2 against face1
if is_edge_contact_geom2:
@@ -2197,34 +2201,29 @@ def _inflate(
@wp.func
def ccd(
def gjk_phase(
# In:
tolerance: float,
cutoff: float,
gjk_iterations: int,
epa_iterations: int,
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
x_1: wp.vec3,
x_2: wp.vec3,
vert: wp.array[wp.vec3],
vert_index: wp.array[int],
face: wp.array[int],
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""General convex collision detection via GJK/EPA."""
) -> Tuple[bool, float, int, wp.vec3, wp.vec3, GJKResult, Geom, Geom]:
"""Run GJK phase of CCD."""
full_margin1 = 0.0
full_margin2 = 0.0
size1 = 0.0
size2 = 0.0
empty = GJKResult()
# determine if the geoms being tested are discrete
is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
# special handling for sphere and capsule (shrink to point and line respectively)
if geomtype1 == GeomType.SPHERE or geomtype1 == GeomType.CAPSULE:
size1 = geom1.size[0]
full_margin1 = size1 + 0.5 * geom1.margin
@@ -2237,7 +2236,6 @@ def ccd(
geom2.margin = 0.0
geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
# special handling for sphere and capsule (shrink to point and line respectively)
if size1 + size2 > 0.0:
cutoff += full_margin1 + full_margin2
result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
@@ -2245,11 +2243,11 @@ def ccd(
# shallow penetration, inflate contact
if result.dist > tolerance:
if result.dist == FLOAT_MAX:
return result.dist, 1, result.x1, result.x2, -1
return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
dist, x1, x2 = _inflate(result, geom1, geom2, geomtype1, geomtype2, full_margin1, full_margin2)
return dist, 1, x1, x2, -1
return False, dist, 1, x1, x2, empty, geom1, geom2
# deep penetration, reset initial conditions and rerun GJK + EPA
# deep penetration: reset initial conditions and rerun GJK + EPA
geom1.margin = full_margin1 - size1
geom1.size = wp.vec3(size1, geom1.size[1], geom1.size[2])
geom2.margin = full_margin2 - size2
@@ -2260,8 +2258,29 @@ def ccd(
# no penetration depth to recover
if result.dist > tolerance or result.dim < 2:
return result.dist, 1, result.x1, result.x2, -1
return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
return True, result.dist, 1, result.x1, result.x2, result, geom1, geom2
@wp.func
def epa_phase(
# In:
tolerance: float,
epa_iterations: int,
result: GJKResult,
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
vert: wp.array[wp.vec3],
vert_index: wp.array[int],
face: wp.array[int],
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
pt = Polytope()
pt.nface = 0
pt.nvert = 0
@@ -2330,6 +2349,7 @@ def ccd(
if pt.status:
return result.dist, 1, result.x1, result.x2, -1
is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
if idx == -1:
return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
@@ -2343,3 +2363,46 @@ def ccd(
idx = -1
return dist, 1, x1, x2, idx
@wp.func
def ccd(
# In:
tolerance: float,
cutoff: float,
gjk_iterations: int,
epa_iterations: int,
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
x_1: wp.vec3,
x_2: wp.vec3,
vert: wp.array[wp.vec3],
vert_index: wp.array[int],
face: wp.array[int],
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""General convex collision detection via GJK/EPA."""
needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
tolerance, cutoff, gjk_iterations, geom1, geom2, geomtype1, geomtype2, x_1, x_2
)
if not needs_epa:
return dist, ncontact, x1, x2, -1
return epa_phase(
tolerance,
epa_iterations,
result,
geom1,
geom2,
geomtype1,
geomtype2,
vert,
vert_index,
face,
face_pr,
face_norm2,
horizon,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+83 -9
View File
@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collisio
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 MJ_MINVAL
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
@@ -183,6 +184,24 @@ def ellipsoid(p: wp.vec3, size: wp.vec3) -> float:
return k0 * (k0 - 1.0) / denom
@wp.func
def capsule(p: wp.vec3, size: wp.vec3) -> float:
r = size[0]
h = size[1]
pz_clamped = wp.clamp(p[2], -h, h)
diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
return wp.length(diff) - r
@wp.func
def cylinder(p: wp.vec3, size: wp.vec3) -> float:
r = size[0]
h = size[1]
dx = wp.length(wp.vec2(p[0], p[1])) - r
dy = wp.abs(p[2]) - h
return wp.min(wp.max(dx, dy), 0.0) + wp.length(wp.vec2(wp.max(dx, 0.0), wp.max(dy, 0.0)))
@wp.func
def grad_sphere(p: wp.vec3) -> wp.vec3:
c = wp.length(p)
@@ -230,6 +249,52 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
return raw_grad / wp.length(raw_grad)
@wp.func
def grad_capsule(p: wp.vec3, size: wp.vec3) -> wp.vec3:
h = size[1]
pz_clamped = wp.clamp(p[2], -h, h)
diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
c = wp.length(diff)
if c > MJ_MINVAL:
return diff / c
else:
return wp.vec3(0.0)
@wp.func
def grad_cylinder(p: wp.vec3, size: wp.vec3) -> wp.vec3:
r = size[0]
h = size[1]
radial_dist = wp.length(wp.vec2(p[0], p[1]))
if radial_dist > MJ_MINVAL:
u = wp.vec3(p[0] / radial_dist, p[1] / radial_dist, 0.0)
else:
u = wp.vec3(0.0)
w = wp.vec3(0.0, 0.0, wp.sign(p[2]))
dx = radial_dist - r
dy = wp.abs(p[2]) - h
if dx > 0.0 and dy > 0.0:
v = wp.vec2(dx, dy)
len_v = wp.length(v)
if len_v > MJ_MINVAL:
return u * (dx / len_v) + w * (dy / len_v)
else:
return wp.vec3(0.0)
elif dx > 0.0:
return u
elif dy > 0.0:
return w
else:
if dx > dy:
return u
else:
return w
@wp.func
def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
"""User-defined SDF function.
@@ -287,16 +352,17 @@ def find_oct(
# check if the node is a leaf
# child indices are relative to root (mesh_octadr offset)
child0 = oct_child[node][0]
# Evaluate this hot leaf predicate eagerly to avoid branch-heavy codegen.
if (
child0 == -1
and oct_child[node][1] == -1
and oct_child[node][2] == -1
and oct_child[node][3] == -1
and oct_child[node][4] == -1
and oct_child[node][5] == -1
and oct_child[node][6] == -1
and oct_child[node][7] == -1
):
int(child0 == -1)
& int(oct_child[node][1] == -1)
& int(oct_child[node][2] == -1)
& int(oct_child[node][3] == -1)
& int(oct_child[node][4] == -1)
& int(oct_child[node][5] == -1)
& int(oct_child[node][6] == -1)
& int(oct_child[node][7] == -1)
) != 0:
for j in range(8):
if not grad:
rx[j] = (
@@ -394,6 +460,10 @@ def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data:
return p[2]
elif type == GeomType.SPHERE:
return sphere(p, attr_vec3)
elif type == GeomType.CAPSULE:
return capsule(p, attr_vec3)
elif type == GeomType.CYLINDER:
return cylinder(p, attr_vec3)
elif type == GeomType.BOX:
return box(p, attr_vec3)
elif type == GeomType.ELLIPSOID:
@@ -452,6 +522,10 @@ def sdf_grad(
return grad
elif type == GeomType.SPHERE:
return grad_sphere(p)
elif type == GeomType.CAPSULE:
return grad_capsule(p, attr_vec3)
elif type == GeomType.CYLINDER:
return grad_cylinder(p, attr_vec3)
elif type == GeomType.BOX:
return grad_box(p, attr_vec3)
elif type == GeomType.ELLIPSOID:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+243 -2
View File
@@ -134,7 +134,10 @@ def _equality_connect(
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
body_invweight0: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
jnt_dofadr: wp.array[int],
dof_bodyid: wp.array[int],
dof_jntid: wp.array[int],
dof_parentid: wp.array[int],
site_bodyid: wp.array[int],
eq_obj1id: wp.array[int],
@@ -154,6 +157,9 @@ def _equality_connect(
site_xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d[wp.spatial_vector],
cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d[wp.spatial_vector],
subtree_linvel_in: wp.array2d[wp.vec3],
njmax_in: int,
njmax_nnz_in: int,
# Data out:
@@ -214,6 +220,7 @@ def _equality_connect(
# compute Jacobian difference (opposite of contact: 0 - 1)
Jqvel = wp.vec3f(0.0, 0.0, 0.0)
Jdotv = wp.vec3f(0.0, 0.0, 0.0)
if is_sparse:
# TODO(team): pre-compute number of non-zeros
@@ -282,6 +289,42 @@ def _equality_connect(
)
j1mj2 = jacp1 - jacp2
jacp1_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
da,
worldid,
)
jacp2_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
da,
worldid,
)
j1mj2_dot = jacp1_dot - jacp2_dot
sparseid0 = rowadr + nnz
sparseid1 = rowadr + rownnz + nnz
sparseid2 = rowadr + 2 * rownnz + nnz
@@ -294,7 +337,9 @@ def _equality_connect(
efc_J_out[worldid, 0, sparseid1] = j1mj2[1]
efc_J_out[worldid, 0, sparseid2] = j1mj2[2]
Jqvel += j1mj2 * qvel_in[worldid, da]
qvel = qvel_in[worldid, da]
Jqvel += j1mj2 * qvel
Jdotv += j1mj2_dot * qvel
nnz += 1
else:
@@ -326,11 +371,49 @@ def _equality_connect(
)
j1mj2 = jacp1 - jacp2
jacp1_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
dofid,
worldid,
)
jacp2_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
dofid,
worldid,
)
j1mj2_dot = jacp1_dot - jacp2_dot
efc_J_out[worldid, efcid0, dofid] = j1mj2[0]
efc_J_out[worldid, efcid1, dofid] = j1mj2[1]
efc_J_out[worldid, efcid2, dofid] = j1mj2[2]
Jqvel += j1mj2 * qvel_in[worldid, dofid]
qvel = qvel_in[worldid, dofid]
Jqvel += j1mj2 * qvel
Jdotv += j1mj2_dot * qvel
body_invweight0_id = worldid % body_invweight0.shape[0]
invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -368,6 +451,8 @@ def _equality_connect(
efc_frictionloss_out,
)
efc_aref_out[worldid, efcidi] -= Jdotv[i]
@wp.kernel
def _equality_joint(
@@ -812,7 +897,10 @@ def _equality_weld(
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
body_invweight0: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
jnt_dofadr: wp.array[int],
dof_bodyid: wp.array[int],
dof_jntid: wp.array[int],
dof_parentid: wp.array[int],
site_bodyid: wp.array[int],
site_quat: wp.array2d[wp.quat],
@@ -834,6 +922,9 @@ def _equality_weld(
site_xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d[wp.spatial_vector],
cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d[wp.spatial_vector],
subtree_linvel_in: wp.array2d[wp.vec3],
njmax_in: int,
njmax_nnz_in: int,
# Data out:
@@ -903,9 +994,43 @@ def _equality_weld(
quat = math.mul_quat(xquat_in[worldid, body1], relpose)
quat1 = math.quat_inv(xquat_in[worldid, body2])
# quat1 = quat_inv(xquat_in[worldid, body2])
q2 = xquat_in[worldid, body2]
quat1 = wp.quat(q2[0], -q2[1], -q2[2], -q2[3])
# compute rotational Jdotv helper quaternions
omega1 = wp.spatial_top(cvel_in[worldid, body1])
omega2 = wp.spatial_top(cvel_in[worldid, body2])
domega = omega1 - omega2
omega1_q = wp.quat(0.0, omega1[0], omega1[1], omega1[2])
omega2_q = wp.quat(0.0, omega2[0], omega2[1], omega2[2])
domega_q = wp.quat(0.0, domega[0], domega[1], domega[2])
if is_site:
qdot0r = math.mul_quat(omega1_q, quat) * 0.5
qfull1 = math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])
qdot1 = math.mul_quat(omega2_q, qfull1) * 0.5
negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
negq1 = wp.quat(qfull1[0], -qfull1[1], -qfull1[2], -qfull1[3])
else:
# qdot0 = mul_quat(xquat_in[worldid, body1], omega1_q) * 0.5
u7 = xquat_in[worldid, body1]
qdot0 = math.mul_quat(omega1_q, xquat_in[worldid, body1]) * 0.5
qdot0r = math.mul_quat(qdot0, relpose)
q1_non_site = xquat_in[worldid, body2]
qdot1 = math.mul_quat(omega2_q, q1_non_site) * 0.5
negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
negq1 = wp.quat(q1_non_site[0], -q1_non_site[1], -q1_non_site[2], -q1_non_site[3])
# compute Jacobian difference (opposite of contact: 0 - 1)
Jqvelp = wp.vec3f(0.0, 0.0, 0.0)
Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
Jdotv_p = wp.vec3f(0.0, 0.0, 0.0)
Jdotv_r0 = wp.vec3f(0.0, 0.0, 0.0)
if is_sparse:
# TODO(team): pre-compute number of non-zeros
@@ -985,6 +1110,44 @@ def _equality_weld(
jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
jacp1_dot, jacr1_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
da,
worldid,
)
jacp2_dot, jacr2_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
da,
worldid,
)
jacdifp_dot = jacp1_dot - jacp2_dot
jacdifr_dot = jacr1_dot - jacr2_dot
sparseid0 = rowadr + nnz
sparseid1 = rowadr + rownnz + nnz
sparseid2 = rowadr + 2 * rownnz + nnz
@@ -1008,6 +1171,8 @@ def _equality_weld(
Jqvelp += jacdifp * qvel_in[worldid, da]
Jqvelr += jacdifr * qvel_in[worldid, da]
Jdotv_p += jacdifp_dot * qvel_in[worldid, da]
Jdotv_r0 += jacdifr_dot * qvel_in[worldid, da]
nnz += 1
else:
@@ -1047,12 +1212,52 @@ def _equality_weld(
jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
jacp1_dot, jacr1_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
dofid,
worldid,
)
jacp2_dot, jacr2_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
dofid,
worldid,
)
jacdifp_dot = jacp1_dot - jacp2_dot
jacdifr_dot = jacr1_dot - jacr2_dot
efc_J_out[worldid, efcid3, dofid] = jacdifr[0]
efc_J_out[worldid, efcid4, dofid] = jacdifr[1]
efc_J_out[worldid, efcid5, dofid] = jacdifr[2]
Jqvelp += jacdifp * qvel_in[worldid, dofid]
Jqvelr += jacdifr * qvel_in[worldid, dofid]
Jdotv_p += jacdifp_dot * qvel_in[worldid, dofid]
Jdotv_r0 += jacdifr_dot * qvel_in[worldid, dofid]
# error is difference in global position and orientation
cpos = pos1 - pos2
@@ -1070,6 +1275,22 @@ def _equality_weld(
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
djrdv_q = wp.quat(0.0, Jdotv_r0[0], Jdotv_r0[1], Jdotv_r0[2])
# Term 1: negqdot1 * domega * q0r
t1a = math.mul_quat(negqdot1, domega_q)
t1 = math.mul_quat(t1a, quat)
# Term 2: negq1 * djrdv * q0r
t2a = math.mul_quat(negq1, djrdv_q)
t2 = math.mul_quat(t2a, quat)
# Term 3: negq1 * domega * qdot0r
t3a = math.mul_quat(negq1, domega_q)
t3 = math.mul_quat(t3a, qdot0r)
Jdotv_r = wp.vec3(t1[1] + t2[1] + t3[1], t1[2] + t2[2] + t3[2], t1[3] + t2[3] + t3[3]) * 0.5 * torquescale
for i in range(3):
_efc_row(
opt_disableflags,
@@ -1096,6 +1317,8 @@ def _equality_weld(
efc_frictionloss_out,
)
efc_aref_out[worldid, efcid + i] -= Jdotv_p[i]
invweight_r = body_invweight0[body_invweight0_id, body1][1] + body_invweight0[body_invweight0_id, body2][1]
for i in range(3):
@@ -1124,6 +1347,8 @@ def _equality_weld(
efc_frictionloss_out,
)
efc_aref_out[worldid, efcid + 3 + i] -= Jdotv_r[i]
@wp.kernel
def _friction_dof(
@@ -1680,6 +1905,7 @@ def _limit_tendon(
)
@cache_kernel
def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
IS_SPARSE = is_sparse
@@ -1803,6 +2029,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
return kernel
@cache_kernel
def _efc_contact_jac_sparse(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -1973,6 +2200,7 @@ def _efc_contact_jac_sparse(cone_type: types.ConeType):
return kernel
@cache_kernel
def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
TILE_SIZE = tile_size
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2117,6 +2345,7 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
return kernel
@cache_kernel
def _efc_contact_update(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2297,7 +2526,10 @@ def make_constraint(m: types.Model, d: types.Data):
m.body_dofnum,
m.body_dofadr,
m.body_invweight0,
m.jnt_type,
m.jnt_dofadr,
m.dof_bodyid,
m.dof_jntid,
m.dof_parentid,
m.site_bodyid,
m.eq_obj1id,
@@ -2316,6 +2548,9 @@ def make_constraint(m: types.Model, d: types.Data):
d.site_xpos,
d.subtree_com,
d.cdof,
d.cvel,
d.cdof_dot,
d.subtree_linvel,
d.njmax,
d.njmax_nnz,
],
@@ -2351,7 +2586,10 @@ def make_constraint(m: types.Model, d: types.Data):
m.body_dofnum,
m.body_dofadr,
m.body_invweight0,
m.jnt_type,
m.jnt_dofadr,
m.dof_bodyid,
m.dof_jntid,
m.dof_parentid,
m.site_bodyid,
m.site_quat,
@@ -2372,6 +2610,9 @@ def make_constraint(m: types.Model, d: types.Data):
d.site_xpos,
d.subtree_com,
d.cdof,
d.cvel,
d.cdof_dot,
d.subtree_linvel,
d.njmax,
d.njmax_nnz,
],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+363 -60
View File
@@ -15,6 +15,7 @@
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src import util_misc
from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
@@ -24,6 +25,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 vec10
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -60,13 +62,47 @@ def _qderiv_actuator_passive_vel(
actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
bias = float(0.0)
if actuator_gaintype[actid] == GainType.AFFINE:
gain = actuator_gainprm[actuator_gainprm_id, actid][2]
elif actuator_gaintype[actid] == GainType.DCMOTOR:
gain = 0.0
dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
gainprm = actuator_gainprm[actuator_gainprm_id, actid]
te = dynprm[0]
# controller velocity derivative: dV/dω
input_mode = int(gainprm[8])
dVdw = 0.0
if input_mode == 1:
dVdw = -gainprm[6] # position: -kd
elif input_mode == 2:
dVdw = -gainprm[4] # velocity: -kp
if te > 0.0:
# stateful current with actearly: d(K*next_act)/dω
# includes both back-EMF (-K) and controller (dVdw) through act_dot
R = wp.max(MJ_MINVAL, gainprm[0])
K = gainprm[1]
s = 1.0 - wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / te)
bias += K * (dVdw - K) * s / R
elif dVdw != 0.0:
# stateless: controller terms only (back-EMF handled in bias block)
R = wp.max(MJ_MINVAL, gainprm[0])
K = gainprm[1]
bias += K * dVdw / R
# LuGre: force includes -sigma1*z_dot, z_dot = a*z + v
# d(sigma1*z_dot)/dv = sigma1*(da/dv*z + 1), ignoring higher-order da/dv*z
sigma1 = dynprm[6]
if sigma1 > 0.0:
bias -= sigma1
else:
gain = 0.0
if actuator_biastype[actid] == BiasType.AFFINE:
bias = actuator_biasprm[actuator_biasprm_id, actid][2]
bias += actuator_biasprm[actuator_biasprm_id, actid][2]
elif actuator_biastype[actid] == BiasType.DCMOTOR:
dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
te = dynprm[0]
@@ -86,11 +122,7 @@ def _qderiv_actuator_passive_vel(
Ta = dynprm[4]
R *= 1.0 + alpha * (T + Ta - T0)
bias = -K * K / wp.max(MJ_MINVAL, R)
else:
bias = 0.0
else:
bias = 0.0
bias += -K * K / wp.max(MJ_MINVAL, R)
if bias == 0.0 and gain == 0.0:
vel_out[worldid, actid] = 0.0
@@ -151,15 +183,15 @@ def _qderiv_actuator_passive_actuation_dense(
actuator_moment_in: wp.array2d[float],
# In:
vel_in: wp.array2d[float],
qMi: wp.array[int],
qMj: wp.array[int],
Mi: wp.array[int],
Mj: wp.array[int],
# Out:
qDeriv_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
dofiid = qMi[elemid]
dofjid = qMj[elemid]
dofiid = Mi[elemid]
dofjid = Mj[elemid]
qderiv_contrib = float(0.0)
for actid in range(nu):
vel = vel_in[worldid, actid]
@@ -195,8 +227,7 @@ def _qderiv_actuator_passive_actuation_dense(
@wp.kernel
def _qderiv_actuator_passive_actuation_sparse(
# Model:
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
M_elemid: wp.array2d[int],
# Data in:
moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d[int],
@@ -204,7 +235,6 @@ def _qderiv_actuator_passive_actuation_sparse(
actuator_moment_in: wp.array2d[float],
# In:
vel_in: wp.array2d[float],
qMj: wp.array[int],
# Out:
qDeriv_out: wp.array3d[float],
):
@@ -231,19 +261,10 @@ def _qderiv_actuator_passive_actuation_sparse(
continue
dofj = moment_colind_in[worldid, rowadrj]
contrib = moment_i * moment_j * vel
# Search the corresponding elemid
# TODO: This could be precalculated for improved performance
row = dofi
col = dofj
row_startk = M_rowadr[row] - 1
row_nnz = M_rownnz[row]
for k in range(row_nnz):
row_startk += 1
if qMj[row_startk] == col:
wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
break
elemid = M_elemid[dofi, dofj]
if elemid >= 0:
contrib = moment_i * moment_j * vel
wp.atomic_add(qDeriv_out[worldid, 0], elemid, contrib)
@wp.kernel
@@ -254,23 +275,29 @@ def _qderiv_actuator_passive(
dof_damping: wp.array2d[float],
dof_dampingpoly: wp.array2d[wp.vec2],
is_sparse: bool,
M_elemid: wp.array2d[int],
# Data in:
qvel_in: wp.array2d[float],
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
qMi: wp.array[int],
qMj: wp.array[int],
Mi: wp.array[int],
Mj: wp.array[int],
qDeriv_in: wp.array3d[float],
# Out:
qDeriv_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
dofiid = qMi[elemid]
dofjid = qMj[elemid]
dofiid = Mi[elemid]
dofjid = Mj[elemid]
madr = M_elemid[dofiid, dofjid]
if is_sparse:
qderiv = qDeriv_in[worldid, 0, elemid]
if madr >= 0:
qderiv = qDeriv_in[worldid, 0, madr]
else:
qderiv = 0.0
else:
qderiv = qDeriv_in[worldid, dofiid, dofjid]
@@ -283,12 +310,13 @@ def _qderiv_actuator_passive(
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
if is_sparse:
qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
if madr >= 0:
qDeriv_out[worldid, 0, madr] = M_in[worldid, 0, madr] - qderiv
else:
qM = qM_in[worldid, dofiid, dofjid] - qderiv
qDeriv_out[worldid, dofiid, dofjid] = qM
M = M_in[worldid, dofiid, dofjid] - qderiv
qDeriv_out[worldid, dofiid, dofjid] = M
if dofiid != dofjid:
qDeriv_out[worldid, dofjid, dofiid] = qM
qDeriv_out[worldid, dofjid, dofiid] = M
# TODO(team): improve performance with tile operations?
@@ -303,18 +331,19 @@ def _qderiv_tendon_damping(
tendon_damping: wp.array2d[float],
tendon_dampingpoly: wp.array2d[wp.vec2],
is_sparse: bool,
M_elemid: wp.array2d[int],
# Data in:
ten_J_in: wp.array2d[float],
ten_velocity_in: wp.array2d[float],
# In:
qMi: wp.array[int],
qMj: wp.array[int],
Mi: wp.array[int],
Mj: wp.array[int],
# Out:
qDeriv_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
dofiid = qMi[elemid]
dofjid = qMj[elemid]
dofiid = Mi[elemid]
dofjid = Mj[elemid]
qderiv = float(0.0)
tendon_damping_id = worldid % tendon_damping.shape[0]
@@ -343,14 +372,283 @@ def _qderiv_tendon_damping(
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
madr = M_elemid[dofiid, dofjid]
if is_sparse:
qDeriv_out[worldid, 0, elemid] -= qderiv
if madr >= 0:
qDeriv_out[worldid, 0, madr] -= qderiv
else:
qDeriv_out[worldid, dofiid, dofjid] -= qderiv
if dofiid != dofjid:
qDeriv_out[worldid, dofjid, dofiid] -= qderiv
@wp.kernel
def deriv_rne_cvel_cdof_dot(
# Model:
body_parentid: wp.array[int],
body_jntnum: wp.array[int],
body_jntadr: wp.array[int],
body_dofadr: wp.array[int],
jnt_type: wp.array[int],
# Data in:
cdof_in: wp.array2d[wp.spatial_vector],
# In:
body_tree_: wp.array[int],
# Out:
Dcvel_out: wp.array3d[wp.spatial_vector],
Dcdof_dot_out: wp.array3d[wp.spatial_vector],
):
"""Forward pass: compute d(cvel)/d(qvel_k) and d(cdof_dot)/d(qvel_k).
Mirrors the accumulation order of comvel for each joint type.
Dcdof_dot for rotation DOFs of free joints (dofid+0..2) is zero because the
forward pass sets cdof_dot[dofid+0..2] = 0. The Dcdof_dot array is
zero-initialized so no explicit write is needed.
"""
worldid, nodeid, dofid = wp.tid()
bodyid = body_tree_[nodeid]
dofadr = body_dofadr[bodyid]
jntid = body_jntadr[bodyid]
jntnum = body_jntnum[bodyid]
pid = body_parentid[bodyid]
cdof = cdof_in[worldid]
# Initialize from parent
cvel_k = Dcvel_out[worldid, pid, dofid]
if jntnum == 0:
Dcvel_out[worldid, bodyid, dofid] = cvel_k
return
dof_i = dofadr
for j in range(jntid, jntid + jntnum):
jnttype = jnt_type[j]
if jnttype == 0: # FREE
# rotation DOFs (dof_i+0..2) contribute to cvel
if dofid >= dof_i and dofid < dof_i + 3:
cvel_k += cdof[dofid]
# cdof_dot for rotation DOFs is zero (set in forward kinematics),
# so Dcdof_dot for rotation DOFs is zero (from wp.zeros init)
# derivative of cdof_dot for translation DOFs 3,4,5
Dcdof_dot_out[worldid, dof_i + 3, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 3])
Dcdof_dot_out[worldid, dof_i + 4, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 4])
Dcdof_dot_out[worldid, dof_i + 5, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 5])
# translation DOFs (dof_i+3..5) contribute to cvel
if dofid >= dof_i + 3 and dofid < dof_i + 6:
cvel_k += cdof[dofid]
dof_i += 6
elif jnttype == 1: # BALL
Dcdof_dot_out[worldid, dof_i + 0, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 0])
Dcdof_dot_out[worldid, dof_i + 1, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 1])
Dcdof_dot_out[worldid, dof_i + 2, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 2])
if dofid >= dof_i and dofid < dof_i + 3:
cvel_k += cdof[dofid]
dof_i += 3
else: # HINGE or SLIDE
Dcdof_dot_out[worldid, dof_i, dofid] = math.motion_cross(cvel_k, cdof[dof_i])
if dofid == dof_i:
cvel_k += cdof[dof_i]
dof_i += 1
Dcvel_out[worldid, bodyid, dofid] = cvel_k
@wp.kernel
def deriv_rne_cacc_cfrcbody_forward(
# Model:
body_parentid: wp.array[int],
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
# Data in:
qvel_in: wp.array2d[float],
cinert_in: wp.array2d[vec10],
cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d[wp.spatial_vector],
# In:
body_tree_: wp.array[int],
Dcvel_in: wp.array3d[wp.spatial_vector],
Dcdof_dot_in: wp.array3d[wp.spatial_vector],
# Out:
Dcacc_out: wp.array3d[wp.spatial_vector],
Dcfrcbody_out: wp.array3d[wp.spatial_vector],
):
"""Forward pass: compute d(cacc)/d(qvel_k) and d(cfrc_body)/d(qvel_k)."""
worldid, nodeid, dofid = wp.tid()
bodyid = body_tree_[nodeid]
dofadr = body_dofadr[bodyid]
dofnum = body_dofnum[bodyid]
pid = body_parentid[bodyid]
qvel = qvel_in[worldid]
dcacc = Dcacc_out[worldid, pid, dofid]
for j in range(dofadr, dofadr + dofnum):
# Term 1: d(cdof_dot * qvel)/d(qvel_k) when j == dofid
if j == dofid:
dcacc += cdof_dot_in[worldid, j]
# Term 2: cdof_dot depends on cvel which depends on qvel_k
dcdofdot = Dcdof_dot_in[worldid, j, dofid]
dcacc += dcdofdot * qvel[j]
Dcacc_out[worldid, bodyid, dofid] = dcacc
# d(cfrc_body)/d(qvel_k)
cinert = cinert_in[worldid, bodyid]
cvel = cvel_in[worldid, bodyid]
dcvel = Dcvel_in[worldid, bodyid, dofid]
# term1 = cinert * d(cacc)/d(qvel_k)
term1 = math.inert_vec(cinert, dcacc)
# term2 = d(cvel x* (cinert * cvel))/d(qvel_k)
cinert_cvel = math.inert_vec(cinert, cvel)
cinert_dcvel = math.inert_vec(cinert, dcvel)
term2 = math.motion_cross_force(dcvel, cinert_cvel) + math.motion_cross_force(cvel, cinert_dcvel)
Dcfrcbody_out[worldid, bodyid, dofid] = term1 + term2
@wp.kernel
def deriv_rne_cfrcbody_backward(
# Model:
body_parentid: wp.array[int],
# In:
body_tree_: wp.array[int],
# Out:
Dcfrcbody_out: wp.array3d[wp.spatial_vector],
):
"""Backward pass: accumulate d(cfrc_body) from children to parents."""
worldid, nodeid, dofid = wp.tid()
bodyid = body_tree_[nodeid]
pid = body_parentid[bodyid]
# body_tree never contains bodyid=0 (worldbody), so pid >= 0 is always valid.
# Siblings at the same level may share a parent; atomic_add handles this.
val = Dcfrcbody_out[worldid, bodyid, dofid]
wp.atomic_add(Dcfrcbody_out[worldid, pid], dofid, val)
@wp.kernel
def deriv_rne_body2jnt_sparse(
# Model:
dof_bodyid: wp.array[int],
# Data in:
cdof_in: wp.array2d[wp.spatial_vector],
# In:
timestep: wp.array[float],
Di: wp.array[int],
Dj: wp.array[int],
Dcfrcbody_in: wp.array3d[wp.spatial_vector],
flg_subtract: bool,
# Out:
qDeriv_out: wp.array3d[float],
):
"""Project body-space RNE derivatives into joint-space qDeriv (sparse)."""
worldid, elemid = wp.tid()
dt = timestep[worldid % timestep.shape[0]]
i = Di[elemid]
j = Dj[elemid]
body_i = dof_bodyid[i]
dcfrc = Dcfrcbody_in[worldid, body_i, j]
term = wp.dot(cdof_in[worldid, i], dcfrc)
if flg_subtract:
wp.atomic_sub(qDeriv_out[worldid, 0], elemid, dt * term)
else:
wp.atomic_add(qDeriv_out[worldid, 0], elemid, dt * term)
def deriv_rne_vel(m: Model, d: Data, out: wp.array3d[float], flg_subtract: bool = False):
"""Compute RNE velocity derivatives and add/subtract from the output.
Implements the analytical derivative of inverse-dynamics Coriolis/centrifugal
forces with respect to joint velocities.
Args:
m: The model (device).
d: The data (device).
out: D-structure output array (nworld, 1, nD) to accumulate RNE terms into.
flg_subtract: If True, subtract the RNE derivatives from output instead of adding them.
"""
# TODO(team): consider caching these allocations
Dcvel = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
Dcdof_dot = wp.zeros((d.nworld, m.nv, m.nv), dtype=wp.spatial_vector)
Dcacc = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
Dcfrcbody = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
# Forward pass 1: compute Dcvel and Dcdof_dot
for body_tree in m.body_tree:
wp.launch(
deriv_rne_cvel_cdof_dot,
dim=(d.nworld, body_tree.size, m.nv),
inputs=[
m.body_parentid,
m.body_jntnum,
m.body_jntadr,
m.body_dofadr,
m.jnt_type,
d.cdof,
body_tree,
],
outputs=[Dcvel, Dcdof_dot],
)
# Forward pass 2: compute Dcacc and Dcfrcbody
for body_tree in m.body_tree:
wp.launch(
deriv_rne_cacc_cfrcbody_forward,
dim=(d.nworld, body_tree.size, m.nv),
inputs=[
m.body_parentid,
m.body_dofnum,
m.body_dofadr,
d.qvel,
d.cinert,
d.cvel,
d.cdof_dot,
body_tree,
Dcvel,
Dcdof_dot,
],
outputs=[Dcacc, Dcfrcbody],
)
# Backward pass: accumulate Dcfrcbody from children to parents
for body_tree in reversed(m.body_tree):
wp.launch(
deriv_rne_cfrcbody_backward,
dim=(d.nworld, body_tree.size, m.nv),
inputs=[m.body_parentid, body_tree],
outputs=[Dcfrcbody],
)
# Project body-space derivatives into joint-space qDeriv (always sparse D-structure)
wp.launch(
deriv_rne_body2jnt_sparse,
dim=(d.nworld, m.qD_fullm_i.size),
inputs=[m.dof_bodyid, d.cdof, m.opt.timestep, m.qD_fullm_i, m.qD_fullm_j, Dcfrcbody, flg_subtract],
outputs=[out],
)
@event_scope
def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
"""Analytical derivative of smooth forces w.r.t. velocities.
@@ -358,12 +656,10 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
Args:
m: The model containing kinematic and dynamic information (device).
d: The data object containing the current state and output arrays (device).
out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
out: M - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
"""
qMi = m.qM_fullm_i
qMj = m.qM_fullm_j
# TODO(team): implicit requires different sparsity structure
Mi = m.M_fullm_i
Mj = m.M_fullm_j
if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
# TODO(team): only clear elements not set by _qderiv_actuator_passive
@@ -399,41 +695,49 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
wp.launch(
_qderiv_actuator_passive_actuation_sparse,
dim=(d.nworld, m.nu),
inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
inputs=[
m.M_elemid,
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind,
d.actuator_moment,
vel,
],
outputs=[out],
)
else:
wp.launch(
_qderiv_actuator_passive_actuation_dense,
dim=(d.nworld, qMi.size),
inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
dim=(d.nworld, Mi.size),
inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, Mi, Mj],
outputs=[out],
)
wp.launch(
_qderiv_actuator_passive,
dim=(d.nworld, qMi.size),
dim=(d.nworld, Mi.size),
inputs=[
m.opt.timestep,
m.opt.disableflags,
m.dof_damping,
m.dof_dampingpoly,
m.is_sparse,
m.M_elemid,
d.qvel,
d.qM,
qMi,
qMj,
d.M,
Mi,
Mj,
out,
],
outputs=[out],
)
else:
# TODO(team): directly utilize qM for these settings
wp.copy(out, d.qM)
# TODO(team): directly utilize M for these settings
wp.copy(out, d.M)
if not (m.opt.disableflags & DisableBit.DAMPER):
wp.launch(
_qderiv_tendon_damping,
dim=(d.nworld, qMi.size),
dim=(d.nworld, Mi.size),
inputs=[
m.ntendon,
m.opt.timestep,
@@ -443,12 +747,11 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
m.tendon_damping,
m.tendon_dampingpoly,
m.is_sparse,
m.M_elemid,
d.ten_J,
d.ten_velocity,
qMi,
qMj,
Mi,
Mj,
],
outputs=[out],
)
# TODO(team): rne derivative
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+82 -25
View File
@@ -20,12 +20,14 @@ 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 history
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
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 types
from mujoco.mjx.third_party.mujoco_warp._src import util_misc
from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
@@ -302,6 +304,9 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
outputs=[d.qpos],
)
# advance history buffers before time advance
history.insert_ctrl_history(m, d)
wp.launch(
_next_time,
dim=d.nworld,
@@ -346,17 +351,18 @@ def _compute_damping_deriv(
def _euler_damp_qfrc_sparse(
# Model:
opt_timestep: wp.array[float],
dof_Madr: wp.array[int],
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
# In:
damp_deriv: wp.array2d[float],
# Out:
qM_integration_out: wp.array3d[float],
M_integration_out: wp.array3d[float],
):
worldid, tid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
adr = dof_Madr[tid]
qM_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
adr = M_rowadr[tid] + M_rownnz[tid] - 1
M_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
@cache_kernel
@@ -366,7 +372,7 @@ def _tile_euler_dense(tile: TileSet):
# Model:
opt_timestep: wp.array[float],
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
efc_Ma_in: wp.array2d[float],
# In:
damp_deriv: wp.array2d[float],
@@ -379,14 +385,14 @@ def _tile_euler_dense(tile: TileSet):
TILE_SIZE = wp.static(tile.size)
dofid = adr_in[nodeid]
M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
damping_tile = wp.tile_load(damp_deriv[worldid], shape=(TILE_SIZE,), offset=(dofid,))
damping_scaled = damping_tile * timestep
qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled)
Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_cholesky(qm_integration_tile)
qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
L_tile = wp.tile_cholesky(qm_integration_tile, fill_mode="upper")
qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile, fill_mode="upper")
wp.tile_store(qacc_out[worldid], qacc_tile, offset=(dofid))
return euler_dense
@@ -409,22 +415,22 @@ def euler(m: Model, d: Data):
)
if m.is_sparse:
qM = wp.clone(d.qM)
M = wp.clone(d.M)
qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
wp.launch(
_euler_damp_qfrc_sparse,
dim=(d.nworld, m.nv),
inputs=[m.opt.timestep, m.dof_Madr, damp_deriv],
outputs=[qM],
inputs=[m.opt.timestep, m.M_rownnz, m.M_rowadr, damp_deriv],
outputs=[M],
)
smooth.factor_solve_i(m, d, qM, qLD, qLDiagInv, qacc, d.efc.Ma)
smooth.factor_solve_i(m, d, M, qLD, qLDiagInv, qacc, d.efc.Ma)
else:
for tile in m.qM_tiles:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_euler_dense(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.opt.timestep, d.qM, d.efc.Ma, damp_deriv, tile.adr],
inputs=[m.opt.timestep, d.M, d.efc.Ma, damp_deriv, tile.adr],
outputs=[qacc],
block_dim=m.block_dim.euler_dense,
)
@@ -573,16 +579,62 @@ def rungekutta4(m: Model, d: Data):
_advance(m, d, qacc_rk, qvel_rk)
@wp.kernel
def _map_m2d(
# Model:
mapM2D: wp.array[int],
is_sparse: bool,
# In:
qDi: wp.array[int],
qDj: wp.array[int],
qH_M: wp.array3d[float],
# Data out:
qLU_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
if is_sparse:
m_idx = mapM2D[elemid]
if m_idx >= 0:
qLU_out[worldid, 0, elemid] = qH_M[worldid, 0, m_idx]
else:
qLU_out[worldid, 0, elemid] = 0.0
else:
i = qDi[elemid]
j = qDj[elemid]
qLU_out[worldid, 0, elemid] = qH_M[worldid, i, j]
@event_scope
def implicit(m: Model, d: Data):
"""Integrates fully implicit in velocity."""
if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
if m.opt.integrator == IntegratorType.IMPLICIT:
qH_M = wp.empty(d.M.shape, dtype=float)
# 1. Compute M - dt * qDeriv_smooth in M-structure
derivative.deriv_smooth_vel(m, d, qH_M)
# 2. Map M-structure to D-structure
wp.launch(
_map_m2d,
dim=(d.nworld, m.nD),
inputs=[m.mapM2D, m.is_sparse, m.qD_fullm_i, m.qD_fullm_j, qH_M],
outputs=[d.qLU],
)
# 3. Compute RNE derivatives, scale by timestep, and subtract in-place from qLU
derivative.deriv_rne_vel(m, d, d.qLU, flg_subtract=True)
# 4. Factorize and solve: qacc = qLU \ Ma
qacc = wp.empty((d.nworld, m.nv), dtype=float)
smooth.factor_solve_lu(m, d, d.qLU, qacc, d.efc.Ma)
_advance(m, d, qacc)
elif ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
if m.is_sparse:
qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
else:
qDeriv = wp.empty(d.qM.shape, dtype=float)
qLD = wp.empty(d.qM.shape, dtype=float)
qDeriv = wp.empty(d.M.shape, dtype=float)
qLD = wp.empty(d.M.shape, dtype=float)
qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
derivative.deriv_smooth_vel(m, d, qDeriv)
qacc = wp.empty((d.nworld, m.nv), dtype=float)
@@ -613,8 +665,7 @@ 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):
if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
island.island(m, d)
smooth.transmission(m, d)
@@ -1098,6 +1149,13 @@ def fwd_actuation(m: Model, d: Data):
d.actuator_force.zero_()
return
# read delayed ctrl (or direct copy if no delay)
if m.nhistory > 0:
ctrl = wp.empty((d.nworld, m.nu), dtype=float)
history.read_ctrl_delayed(m, d, ctrl)
else:
ctrl = d.ctrl
wp.launch(
_actuator_force,
dim=(d.nworld, m.nu),
@@ -1122,7 +1180,7 @@ def fwd_actuation(m: Model, d: Data):
m.actuator_acc0,
m.actuator_lengthrange,
d.act,
d.ctrl,
ctrl,
d.actuator_length,
d.actuator_velocity,
m.opt.disableflags & DisableBit.CLAMPCTRL,
@@ -1221,7 +1279,7 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
xfrc_accumulate(m, d, d.qfrc_smooth)
if factorize:
smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
else:
smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
@@ -1266,7 +1324,7 @@ def step(m: Model, d: Data):
euler(m, d)
elif m.opt.integrator == IntegratorType.RK4:
rungekutta4(m, d)
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
elif m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
implicit(m, d)
else:
raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
@@ -1307,8 +1365,7 @@ def step2(m: Model, d: Data):
sensor.sensor_acc(m, d)
# integrate with Euler or implicitfast
# TODO(team): implicit
if m.opt.integrator == IntegratorType.IMPLICITFAST:
if m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
implicit(m, d)
else:
# note: RK4 defaults to Euler
mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
Vendored
+925
View File
@@ -0,0 +1,925 @@
# 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.
# ==============================================================================
import warp as wp
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 Data
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
wp.set_module_options({"enable_backward": False})
@wp.func
def _history_physical_index(cursor: int, n: int, logical: int) -> int:
"""Convert logical index (0=oldest, n-1=newest) to physical index."""
return (cursor + 1 + logical) % n
@wp.func
def _history_find_index(
# In:
buf: wp.array2d[float],
worldid: int,
buf_offset: int,
n: int,
cursor: int,
t: float,
) -> int:
"""Find logical index i such that times[i-1] < t <= times[i].
Returns 0 if t <= times[oldest], n if t > times[newest].
Uses circular binary search matching MuJoCo C historyFindIndex.
"""
times_offset = buf_offset + 2
oldest_phys = _history_physical_index(cursor, n, 0)
newest_phys = _history_physical_index(cursor, n, n - 1)
t_oldest = buf[worldid, times_offset + oldest_phys]
t_newest = buf[worldid, times_offset + newest_phys]
# before or at first element
if t <= t_oldest:
return 0
# after last element
if t > t_newest:
return n
# circular binary search: find smallest logical i such that times[phys(i)] >= t
lo = int(0)
hi = int(n - 1)
while hi - lo > 1:
mid = int((lo + hi) >> 1)
mid_phys = _history_physical_index(cursor, n, mid)
if buf[worldid, times_offset + mid_phys] < t:
lo = mid
else:
hi = mid
return hi
@wp.func
def _history_read_scalar(
# In:
buf: wp.array2d[float],
worldid: int,
buf_offset: int,
n: int,
t: float,
interp: int,
) -> float:
"""Read a scalar value from history buffer at time t.
interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
"""
cursor = int(buf[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
oldest_phys = _history_physical_index(cursor, n, 0)
newest_phys = _history_physical_index(cursor, n, n - 1)
t_oldest = buf[worldid, times_offset + oldest_phys]
t_newest = buf[worldid, times_offset + newest_phys]
# extrapolate before oldest
if t <= t_oldest + 1e-6:
return buf[worldid, values_offset + oldest_phys]
# extrapolate after newest
if t >= t_newest - 1e-6:
return buf[worldid, values_offset + newest_phys]
# find bracketing index
i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
phys_i = _history_physical_index(cursor, n, i)
# exact match
if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
return buf[worldid, values_offset + phys_i]
phys_lo = _history_physical_index(cursor, n, i - 1)
phys_hi = phys_i
# zero-order hold
if interp == 0:
return buf[worldid, values_offset + phys_lo]
dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
v_lo = buf[worldid, values_offset + phys_lo]
v_hi = buf[worldid, values_offset + phys_hi]
# linear interpolation
if interp == 1:
return v_lo + alpha * (v_hi - v_lo)
# cubic spline interpolation (Catmull-Rom)
alpha2 = alpha * alpha
alpha3 = alpha2 * alpha
h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
h10 = alpha3 - 2.0 * alpha2 + alpha
h01 = -2.0 * alpha3 + 3.0 * alpha2
h11 = alpha3 - alpha2
# finite-differenced Catmull-Rom slopes, 0 at endpoints
m_lo = 0.0
if i > 1:
phys_lo_prev = _history_physical_index(cursor, n, i - 2)
dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev]) / dt_lo
m_hi = 0.0
if i < n - 1:
phys_hi_next = _history_physical_index(cursor, n, i + 1)
dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
m_hi = (buf[worldid, values_offset + phys_hi_next] - v_lo) / dt_hi
return h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
@wp.func
def _history_read_vector(
# In:
adr: int,
buf: wp.array2d[float],
worldid: int,
buf_offset: int,
n: int,
dim: int,
t: float,
interp: int,
# Data out:
sensordata_out: wp.array2d[float],
) -> int:
"""Read a vector value from history buffer at time t into sensordata.
Returns 1 on success (value written to sensordata).
interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
"""
cursor = int(buf[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
oldest_phys = _history_physical_index(cursor, n, 0)
newest_phys = _history_physical_index(cursor, n, n - 1)
t_oldest = buf[worldid, times_offset + oldest_phys]
t_newest = buf[worldid, times_offset + newest_phys]
# extrapolate before oldest: copy oldest
if t <= t_oldest + 1e-6:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + oldest_phys * dim + d]
return 1
# extrapolate after newest: copy newest
if t >= t_newest - 1e-6:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + newest_phys * dim + d]
return 1
# find bracketing index
i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
phys_i = _history_physical_index(cursor, n, i)
# exact match
if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_i * dim + d]
return 1
phys_lo = _history_physical_index(cursor, n, i - 1)
phys_hi = phys_i
# zero-order hold
if interp == 0:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_lo * dim + d]
return 1
dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
# linear interpolation
if interp == 1:
for d in range(dim):
v_lo = buf[worldid, values_offset + phys_lo * dim + d]
v_hi = buf[worldid, values_offset + phys_hi * dim + d]
sensordata_out[worldid, adr + d] = v_lo + alpha * (v_hi - v_lo)
return 1
# cubic spline interpolation (Catmull-Rom)
alpha2 = alpha * alpha
alpha3 = alpha2 * alpha
h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
h10 = alpha3 - 2.0 * alpha2 + alpha
h01 = -2.0 * alpha3 + 3.0 * alpha2
h11 = alpha3 - alpha2
for d in range(dim):
v_lo = buf[worldid, values_offset + phys_lo * dim + d]
v_hi = buf[worldid, values_offset + phys_hi * dim + d]
# finite-differenced Catmull-Rom slopes, 0 at endpoints
m_lo = 0.0
if i > 1:
phys_lo_prev = _history_physical_index(cursor, n, i - 2)
dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev * dim + d]) / dt_lo
m_hi = 0.0
if i < n - 1:
phys_hi_next = _history_physical_index(cursor, n, i + 1)
dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
m_hi = (buf[worldid, values_offset + phys_hi_next * dim + d] - v_lo) / dt_hi
sensordata_out[worldid, adr + d] = h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
return 1
@wp.func
def _history_insert_scalar(
# In:
worldid: int,
buf_offset: int,
n: int,
t: float,
value: float,
# Out:
buf_out: wp.array2d[float],
):
"""Insert a scalar value into history buffer at time t."""
cursor = int(buf_out[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
# exact match
if i < n:
phys_i = _history_physical_index(cursor, n, i)
if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
buf_out[worldid, values_offset + phys_i] = value
return
# older than oldest: replace oldest
if i == 0:
oldest_phys = _history_physical_index(cursor, n, 0)
buf_out[worldid, times_offset + oldest_phys] = t
buf_out[worldid, values_offset + oldest_phys] = value
return
# newer than newest: advance cursor
if i == n:
cursor = (cursor + 1) % n
buf_out[worldid, buf_offset + 1] = float(cursor)
buf_out[worldid, times_offset + cursor] = t
buf_out[worldid, values_offset + cursor] = value
return
# out-of-order: shift [1, i-1] left, insert at i-1
for j in range(i - 1):
src_phys = _history_physical_index(cursor, n, j + 1)
dst_phys = _history_physical_index(cursor, n, j)
buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
buf_out[worldid, values_offset + dst_phys] = buf_out[worldid, values_offset + src_phys]
insert_phys = _history_physical_index(cursor, n, i - 1)
buf_out[worldid, times_offset + insert_phys] = t
buf_out[worldid, values_offset + insert_phys] = value
@wp.func
def _history_insert_vector(
# In:
worldid: int,
buf_offset: int,
n: int,
dim: int,
t: float,
src: wp.array2d[float],
src_adr: int,
# Out:
buf_out: wp.array2d[float],
):
"""Insert a vector value from src[worldid, src_adr:src_adr+dim] into history buffer at time t."""
cursor = int(buf_out[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
slot_phys = -1
# exact match
if i < n:
phys_i = _history_physical_index(cursor, n, i)
if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
slot_phys = phys_i
if slot_phys < 0:
if i == 0:
# older than oldest: replace oldest
slot_phys = _history_physical_index(cursor, n, 0)
buf_out[worldid, times_offset + slot_phys] = t
elif i == n:
# newer than newest: advance cursor
cursor = (cursor + 1) % n
buf_out[worldid, buf_offset + 1] = float(cursor)
slot_phys = cursor
buf_out[worldid, times_offset + slot_phys] = t
else:
# out-of-order: shift [1, i-1] left, insert at i-1
for j in range(i - 1):
src_phys = _history_physical_index(cursor, n, j + 1)
dst_phys = _history_physical_index(cursor, n, j)
buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
for d in range(dim):
buf_out[worldid, values_offset + dst_phys * dim + d] = buf_out[worldid, values_offset + src_phys * dim + d]
slot_phys = _history_physical_index(cursor, n, i - 1)
buf_out[worldid, times_offset + slot_phys] = t
# copy values
for d in range(dim):
buf_out[worldid, values_offset + slot_phys * dim + d] = src[worldid, src_adr + d]
@wp.kernel
def _read_ctrl_delayed_kernel(
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
actuator_delay: wp.array[float],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
# Data out:
ctrl_out: wp.array2d[float],
):
"""Read delayed ctrl for each actuator."""
worldid, uid = wp.tid()
hist = actuator_history[uid]
nsample = hist[0]
delay = actuator_delay[uid]
if nsample == 0 or delay == 0.0:
# no delay: direct copy
ctrl_out[worldid, uid] = ctrl_in[worldid, uid]
else:
interp = hist[1]
buf_offset = actuator_historyadr[uid]
t = time_in[worldid] - delay
ctrl_out[worldid, uid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp)
@wp.kernel
def _insert_ctrl_history_kernel(
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
# Data in:
time_in: wp.array[float],
ctrl_in: wp.array2d[float],
# Data out:
history_out: wp.array2d[float],
):
"""Insert current ctrl into history buffers."""
worldid, uid = wp.tid()
hist = actuator_history[uid]
nsample = hist[0]
if nsample == 0:
return
buf_offset = actuator_historyadr[uid]
t = time_in[worldid]
value = ctrl_in[worldid, uid]
_history_insert_scalar(worldid, buf_offset, nsample, t, value, history_out)
@wp.kernel
def _insert_sensor_history_stage(
# Model:
sensor_dim: wp.array[int],
sensor_adr: wp.array[int],
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_delay: wp.array[float],
sensor_interval: wp.array[wp.vec2],
# Data in:
time_in: wp.array[float],
sensordata_in: wp.array2d[float],
# In:
sensor_ids: wp.array[int],
# Data out:
history_out: wp.array2d[float],
):
"""Insert current sensor values into history buffers for specific sensor IDs."""
worldid, idx = wp.tid()
sid = sensor_ids[idx]
hist = sensor_history[sid]
nsample = hist[0]
if nsample == 0:
return
buf_offset = sensor_historyadr[sid]
dim = sensor_dim[sid]
interval_val = sensor_interval[sid]
period = interval_val[0]
t = time_in[worldid]
if period > 0.0:
# interval mode: check if condition is satisfied
time_prev = history_out[worldid, buf_offset] # user slot stores time_prev
if time_prev + period <= t:
# advance time_prev by exact period
history_out[worldid, buf_offset] = time_prev + period
# insert sensor value
_history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
else:
_history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
@wp.kernel
def _apply_sensor_delay_kernel(
# Model:
sensor_dim: wp.array[int],
sensor_adr: wp.array[int],
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_delay: wp.array[float],
sensor_interval: wp.array[wp.vec2],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
# In:
sensor_ids: wp.array[int],
# Data out:
sensordata_out: wp.array2d[float],
):
"""Apply delay/interval logic for sensors after computation.
TODO(team): Revisit always-compute decision for computationally expensive sensors
with interval/period (e.g., raytracers)
"""
worldid, idx = wp.tid()
sid = sensor_ids[idx]
hist = sensor_history[sid]
nsample = hist[0]
if nsample <= 0:
return
delay = sensor_delay[sid]
dim = sensor_dim[sid]
interp = hist[1]
buf_offset = sensor_historyadr[sid]
t = time_in[worldid]
if delay > 0.0:
# delay > 0: read delayed value from buffer
_history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t - delay, interp, sensordata_out)
else:
# interval-only (delay == 0, interval > 0): check interval condition
interval_val = sensor_interval[sid]
period = interval_val[0]
if period > 0.0:
time_prev = history_in[worldid, buf_offset] # user slot
if time_prev + period > t:
# interval condition not satisfied: read from buffer
_history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t, interp, sensordata_out)
# else: interval condition satisfied, keep computed value
def read_ctrl_delayed(m: Model, d: Data, ctrl: wp.array2d[float]):
"""Read delayed ctrl values for all actuators."""
if m.nhistory == 0:
wp.copy(ctrl, d.ctrl)
return
wp.launch(
_read_ctrl_delayed_kernel,
dim=(d.nworld, m.nu),
inputs=[
m.actuator_history,
m.actuator_historyadr,
m.actuator_delay,
d.time,
d.history,
d.ctrl,
],
outputs=[ctrl],
)
def insert_ctrl_history(m: Model, d: Data):
"""Insert current ctrl values into history buffers."""
if m.nhistory == 0 or m.nu == 0:
return
wp.launch(
_insert_ctrl_history_kernel,
dim=(d.nworld, m.nu),
inputs=[
m.actuator_history,
m.actuator_historyadr,
d.time,
d.ctrl,
],
outputs=[d.history],
)
def apply_sensor_delay(m: Model, d: Data, sensorid: wp.array[int]):
"""Apply delay/interval logic for given sensors after computation.
Matches MuJoCo C architecture where the delayed read (mj_sensorPos) occurs
before the fresh value insert (mj_advance). We save fresh sensordata,
overwrite with delayed values, then insert the saved fresh values.
"""
if m.nhistory == 0 or sensorid.shape[0] == 0:
return
# Save fresh sensordata before delay overwrite
fresh_sensordata = wp.empty_like(d.sensordata)
wp.copy(fresh_sensordata, d.sensordata)
# Read delayed values from buffer → overwrite sensordata
wp.launch(
_apply_sensor_delay_kernel,
dim=(d.nworld, sensorid.shape[0]),
inputs=[
m.sensor_dim,
m.sensor_adr,
m.sensor_history,
m.sensor_historyadr,
m.sensor_delay,
m.sensor_interval,
d.time,
d.history,
sensorid,
],
outputs=[d.sensordata],
)
# Insert saved fresh sensor values into history buffers
wp.launch(
_insert_sensor_history_stage,
dim=(d.nworld, sensorid.shape[0]),
inputs=[
m.sensor_dim,
m.sensor_adr,
m.sensor_history,
m.sensor_historyadr,
m.sensor_delay,
m.sensor_interval,
d.time,
fresh_sensordata,
sensorid,
],
outputs=[d.history],
)
@wp.kernel
def _read_ctrl_kernel(
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
actuator_delay: wp.array[float],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
# In:
uid: int,
interp: int,
# Out:
result_out: wp.array[float],
):
"""Read delayed ctrl for 1 actuator across all worlds."""
worldid = wp.tid()
hist = actuator_history[uid]
nsample = hist[0]
if nsample == 0:
result_out[worldid] = ctrl_in[worldid, uid]
else:
interp_val = interp
if interp_val < 0:
interp_val = hist[1]
delay = actuator_delay[uid]
buf_offset = actuator_historyadr[uid]
t = time_in[worldid] - delay
result_out[worldid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp_val)
def read_ctrl(
m: Model,
d: Data,
ctrlid: int,
time: wp.array[float],
interp: int,
result: wp.array2d[float],
):
"""Read delayed ctrl for 1 actuator across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
ctrlid: actuator index.
time: query time per world (nworld,).
interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
result: output buffer (nworld,).
"""
wp.launch(
_read_ctrl_kernel,
dim=(d.nworld,),
inputs=[
m.actuator_history,
m.actuator_historyadr,
m.actuator_delay,
time,
d.history,
d.ctrl,
ctrlid,
interp,
],
outputs=[result],
)
@wp.kernel
def _read_sensor_kernel(
# Model:
sensor_dim: wp.array[int],
sensor_adr: wp.array[int],
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_delay: wp.array[float],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
sensordata_in: wp.array2d[float],
# In:
sid: int,
interp: int,
# Out:
result_out: wp.array2d[float],
):
"""Read delayed sensor for 1 sensor across all worlds."""
worldid = wp.tid()
hist = sensor_history[sid]
nsample = hist[0]
dim = sensor_dim[sid]
adr = sensor_adr[sid]
if nsample == 0:
for i in range(dim):
result_out[worldid, i] = sensordata_in[worldid, adr + i]
else:
interp_val = interp
if interp_val < 0:
interp_val = hist[1]
delay = sensor_delay[sid]
buf_offset = sensor_historyadr[sid]
t = time_in[worldid] - delay
_history_read_vector(
0, # write to result_out starting at index 0 (not global sensor adr)
history_in,
worldid,
buf_offset,
nsample,
dim,
t,
interp_val,
result_out,
)
def read_sensor(
m: Model,
d: Data,
sensorid: int,
time: wp.array[float],
interp: int,
result: wp.array2d[float],
):
"""Read delayed sensor for 1 sensor across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
sensorid: sensor index.
time: query time per world (nworld,).
interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
result: output buffer (nworld, dim).
"""
wp.launch(
_read_sensor_kernel,
dim=(d.nworld,),
inputs=[
m.sensor_dim,
m.sensor_adr,
m.sensor_history,
m.sensor_historyadr,
m.sensor_delay,
time,
d.history,
d.sensordata,
sensorid,
interp,
],
outputs=[result],
)
@wp.kernel
def _init_ctrl_history_kernel(
# kernel_analyzer: off
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
# In:
ctrlid: int,
times: wp.array[float],
values: wp.array2d[float],
has_times: int,
# Data out:
history_out: wp.array2d[float],
# kernel_analyzer: on
):
"""Initialize history buffer for 1 actuator across all worlds."""
worldid = wp.tid()
nsample = actuator_history[ctrlid][0]
buf_offset = actuator_historyadr[ctrlid]
# preserve user slot
user = history_out[worldid, buf_offset]
# cursor = 0 (samples in order, newest at index nsample-1)
history_out[worldid, buf_offset + 1] = float(nsample - 1)
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + nsample
for i in range(nsample):
if has_times != 0:
history_out[worldid, times_offset + i] = times[i]
else:
history_out[worldid, times_offset + i] = -MJ_MAXVAL
history_out[worldid, values_offset + i] = values[worldid, i]
# restore user slot
history_out[worldid, buf_offset] = user
def init_ctrl_history(
m: Model,
d: Data,
ctrlid: int,
times: wp.array[float],
values: wp.array2d[float],
):
"""Initialize history buffer for 1 actuator across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
ctrlid: actuator index.
times: timestamps or None (nsample,).
values: ctrl values (nworld, nsample).
Raises:
ValueError: If times are not strictly increasing.
"""
has_times = 0 if times is None else 1
if times is not None:
t_np = times.numpy()
for i in range(len(t_np) - 1):
if t_np[i + 1] - t_np[i] < MJ_MINVAL:
raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
if times is None:
times = wp.empty(0, dtype=float)
wp.launch(
_init_ctrl_history_kernel,
dim=(d.nworld,),
inputs=[
m.actuator_history,
m.actuator_historyadr,
ctrlid,
times,
values,
has_times,
],
outputs=[d.history],
)
# kernel_analyzer: off
@wp.kernel
def _init_sensor_history_kernel(
# Model:
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_dim_arr: wp.array[int],
# In:
sensorid: int,
times: wp.array[float],
values: wp.array2d[float],
phase: wp.array[float],
has_times: int,
# Data out:
history_out: wp.array2d[float],
):
# kernel_analyzer: on
"""Initialize history buffer for 1 sensor across all worlds."""
worldid = wp.tid()
nsample = sensor_history[sensorid][0]
dim = sensor_dim_arr[sensorid]
buf_offset = sensor_historyadr[sensorid]
# set user slot (phase = last computation time for interval sensors)
history_out[worldid, buf_offset] = phase[worldid]
# cursor = 0 (samples in order, newest at index nsample-1)
history_out[worldid, buf_offset + 1] = float(nsample - 1)
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + nsample
for i in range(nsample):
if has_times != 0:
history_out[worldid, times_offset + i] = times[i]
else:
history_out[worldid, times_offset + i] = -MJ_MAXVAL
for j in range(dim):
history_out[worldid, values_offset + i * dim + j] = values[worldid, i * dim + j]
def init_sensor_history(
m: Model,
d: Data,
sensorid: int,
times: wp.array[float],
values: wp.array2d[float],
phase: wp.array[float],
):
"""Initialize history buffer for 1 sensor across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
sensorid: sensor index.
times: timestamps or None (nsample,).
values: sensor values (nworld, nsample * dim).
phase: user slot value per world (nworld,).
Raises:
ValueError: If times are not strictly increasing.
"""
has_times = 0 if times is None else 1
if times is not None:
t_np = times.numpy()
for i in range(len(t_np) - 1):
if t_np[i + 1] - t_np[i] < MJ_MINVAL:
raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
if times is None:
times = wp.empty(0, dtype=float)
wp.launch(
_init_sensor_history_kernel,
dim=(d.nworld,),
inputs=[
m.sensor_history,
m.sensor_historyadr,
m.sensor_dim,
sensorid,
times,
values,
phase,
has_times,
],
outputs=[d.history],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
Vendored
+5 -5
View File
@@ -95,9 +95,9 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
# TODO(team): qacc = d.qacc if (m.dof_damping == 0.0).all()
# set qfrc = (d.qM + m.opt.timestep * diag(m.dof_damping)) * d.qacc
# set qfrc = (d.M + m.opt.timestep * diag(m.dof_damping)) * d.qacc
# d.qM @ d.qacc
# d.M @ d.qacc
support.mul_m(m, d, qfrc, d.qacc)
# qfrc += m.opt.timestep * damp_deriv * d.qacc
@@ -109,16 +109,16 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
)
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
if m.is_sparse:
qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
else:
qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
derivative.deriv_smooth_vel(m, d, qDeriv)
mul_m(m, d, qfrc, d.qacc, M=qDeriv)
smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, qacc, qfrc)
smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, qacc, qfrc)
else:
raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
# solve for qacc: qfrc = d.qM @ d.qacc
# solve for qacc: qfrc = d.M @ d.qacc
smooth.solve_m(m, d, qacc, qfrc)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+330 -102
View File
@@ -33,6 +33,9 @@ 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
# TODO(team): remove after improving island solver performance
ENABLE_ISLANDS = False
def _is_array_spec(typ) -> bool:
"""Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
@@ -64,6 +67,61 @@ def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
return array
def _create_constraint(
mjm,
nworld: int,
njmax: int,
njmax_nnz: int,
sizes: dict,
island_enabled: bool,
mjd=None,
) -> types.Constraint:
"""Construct a types.Constraint with standard and island local fields allocated properly."""
efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
sparse = is_sparse(mjm)
for f in dataclasses.fields(types.Constraint):
if f.name == "itype":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
elif f.name == "iid":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
elif f.name == "iJ_rownnz":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
elif f.name == "iJ_rowadr":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
elif f.name == "iJ_colind":
efc_kwargs[f.name] = wp.empty((nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, 0, 0), dtype=int)
elif f.name == "iJ":
efc_kwargs[f.name] = wp.empty(
(nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, njmax if island_enabled else 0, mjm.nv),
dtype=float,
)
elif f.name == "iD":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "iaref":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "ifrictionloss":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "iforce":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "istate":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
else:
if f.name in efc_kwargs:
continue
if mjd is not None:
shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
val = np.zeros(shape, dtype=f.type.dtype)
if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
else:
efc_kwargs[f.name] = _create_array(None, f.type, sizes)
return types.Constraint(**efc_kwargs)
def is_sparse(mjm: mujoco.MjModel) -> bool:
if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
if mjm.nv > 32:
@@ -183,6 +241,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL))
opt = types.Option(**opt_kwargs)
# islands are disabled by default while performance is being improved
# override by setting io.ENABLE_ISLANDS = True
# TODO(team): remove after improving island solver performance
if not ENABLE_ISLANDS:
opt.disableflags |= types.DisableBit.ISLAND
# C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU
# adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly
# bouncing around the optimal solution
@@ -271,7 +335,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
jnt_limited_slide_hinge = mjm.jnt_limited & np.isin(mjm.jnt_type, (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE))
m.jnt_limited_slide_hinge_adr = np.nonzero(jnt_limited_slide_hinge)[0]
m.jnt_limited_ball_adr = np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0]
m.dof_tri_row, m.dof_tri_col = np.tril_indices(mjm.nv)
m.dof_tri_row, m.dof_tri_col = np.triu_indices(mjm.nv)
# precompute body_isdofancestor: which DOFs affect each body
# TODO: Investigate alternative approach such as bitmap
@@ -585,14 +649,14 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]])
]
# qM_tiles records the block diagonal structure of qM
# M_tiles records the block diagonal structure of M
tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1]
tiles = {}
for i in range(len(tile_corners)):
tile_beg = tile_corners[i]
tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
m.qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
# qLD_updates has dof tree ordering of qLD updates for sparse factor m
qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
@@ -620,41 +684,69 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
m.qLD_level_offsets = level_offsets
# indices for sparse qM_fullm (used in solver)
m.qM_fullm_i, m.qM_fullm_j = [], []
# Indices for sparse M_fullm (used in solver). M_fullm_i/j are built by
# walking dof_parentid for each dof, so for joint types whose internal block
# MuJoCo stores diagonal-only in the compact (M_rownnz, M_rowadr) layout
# (e.g. free joints), the chain-aware layout here has more entries per row
# than the compact layout.
m.M_fullm_i, m.M_fullm_j = [], []
for i in range(mjm.nv):
j = i
while j > -1:
m.qM_fullm_i.append(i)
m.qM_fullm_j.append(j)
m.M_fullm_i.append(i)
m.M_fullm_j.append(j)
j = mjm.dof_parentid[j]
# M_elemid maps (row, col) -> madr index in the native CSR M layout
M_elemid = np.full((mjm.nv, mjm.nv), -1, dtype=np.int32)
for i in range(mjm.nv):
rowadr = mjm.M_rowadr[i]
rownnz = mjm.M_rownnz[i]
for k in range(rownnz):
madr = rowadr + k
col = int(mjm.M_colind[madr])
M_elemid[i, col] = madr
m.M_elemid = M_elemid
upper_j, upper_i = np.triu_indices(mjm.nv)
upper_elemid = M_elemid[upper_i, upper_j]
valid_mask = upper_elemid != -1
m.M_fullm_upper_i = upper_j[valid_mask].tolist()
m.M_fullm_upper_j = upper_i[valid_mask].tolist()
m.M_fullm_upper_elemid = upper_elemid[valid_mask].tolist()
# indices for sparse qD_fullm (used in RNE derivatives)
# D-structure is the full square sparsity pattern (both upper and lower triangle)
m.qD_fullm_i, m.qD_fullm_j = [], []
for i in range(mjm.nv):
rowadr = mjm.D_rowadr[i]
rownnz = mjm.D_rownnz[i]
for k in range(rownnz):
m.qD_fullm_i.append(i)
m.qD_fullm_j.append(int(mjm.D_colind[rowadr + k]))
m.nD = mjm.nD
# Gather-based sparse mul_m: for each row, all (col, madr) including diagonal
row_elements = [[] for _ in range(mjm.nv)]
# Add diagonal
for i in range(mjm.nv):
row_elements[i].append((i, mjm.dof_Madr[i]))
# Add off-diagonals: ancestors (lower) and descendants (upper)
for i in range(mjm.nv):
madr_ij, j = mjm.dof_Madr[i], i
while True:
madr_ij, j = madr_ij + 1, mjm.dof_parentid[j]
if j == -1:
break
row_elements[i].append((j, madr_ij)) # row i gathers M[i,j] * vec[j]
row_elements[j].append((i, madr_ij)) # row j gathers M[j,i] * vec[i]
rowadr = mjm.M_rowadr[i]
rownnz = mjm.M_rownnz[i]
for k in range(rownnz):
madr = rowadr + k
col = int(mjm.M_colind[madr])
row_elements[i].append((col, madr)) # row i gathers M[i,col] * vec[col]
if i != col:
row_elements[col].append((i, madr)) # row col gathers M[i,col] * vec[i]
# Flatten into CSR-like arrays
m.qM_mulm_rowadr = [0]
m.qM_mulm_col = []
m.qM_mulm_madr = []
m.M_mulm_rowadr = [0]
m.M_mulm_col = []
m.M_mulm_madr = []
for i in range(mjm.nv):
for col, madr in row_elements[i]:
m.qM_mulm_col.append(col)
m.qM_mulm_madr.append(madr)
m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
m.M_mulm_col.append(col)
m.M_mulm_madr.append(madr)
m.M_mulm_rowadr.append(len(m.M_mulm_col))
m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
@@ -886,6 +978,42 @@ def _resolve_batch_size(na: int | None, n: int | None, nworld: int, default: int
return default
def _allocate_island_arrays(
mjm: mujoco.MjModel,
d: types.Data,
nworld: int,
njmax: int,
island_enabled: bool,
mjd: mujoco.MjData,
):
ntree_size = mjm.ntree if island_enabled else 0
nv_size = mjm.nv if island_enabled else 0
njmax_size = njmax if island_enabled else 0
d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1 if island_enabled else 0)), dtype=int)
d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
d.island_nf = wp.empty((nworld, ntree_size), dtype=int)
d.island_efcadr = wp.empty((nworld, ntree_size), dtype=int)
d.nidof = wp.empty((nworld if island_enabled else 0,), dtype=int)
d.map_dof2idof = wp.empty((nworld, nv_size), dtype=int)
d.map_idof2dof = wp.empty((nworld, nv_size), dtype=int)
d.map_efc2iefc = wp.empty((nworld, njmax_size), dtype=int)
d.map_iefc2efc = wp.empty((nworld, njmax_size), dtype=int)
d.dof_islandid = wp.empty((nworld, nv_size), dtype=int)
d.efc_islandid = wp.empty((nworld, njmax_size), dtype=int)
d.iqacc = wp.empty((nworld, nv_size), dtype=float)
d.iqacc_smooth = wp.empty((nworld, nv_size), dtype=float)
d.iqfrc_smooth = wp.empty((nworld, nv_size), dtype=float)
d.iqfrc_constraint = wp.empty((nworld, nv_size), dtype=float)
def make_data(
mjm: mujoco.MjModel,
nworld: int = 1,
@@ -967,7 +1095,8 @@ def make_data(
contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS)
if is_sparse(mjm):
efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
@@ -1005,7 +1134,7 @@ def make_data(
"njmax": njmax,
"njmax_pad": sizes["njmax_pad"],
"njmax_nnz": njmax_nnz,
"qM": None,
"M": None,
"qLD": None,
# world body
"xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
@@ -1024,6 +1153,24 @@ def make_data(
# island arrays
"nisland": None,
"tree_island": None,
"dof_island": None,
"island_dofadr": None,
"island_nv": None,
"island_nefc": None,
"island_ne": None,
"island_nf": None,
"island_efcadr": None,
"nidof": None,
"map_dof2idof": None,
"map_idof2dof": None,
"map_efc2iefc": None,
"map_iefc2efc": None,
"dof_islandid": None,
"efc_islandid": None,
"iqacc": None,
"iqacc_smooth": None,
"iqfrc_smooth": None,
"iqfrc_constraint": None,
}
for f in dataclasses.fields(types.Data):
if f.name in d_kwargs:
@@ -1033,15 +1180,13 @@ def make_data(
d = types.Data(**d_kwargs)
if is_sparse(mjm):
d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
d.M = wp.zeros((nworld, 1, mjm.nC), dtype=float)
d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
else:
d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
d.M = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
# island discovery arrays
d.nisland = wp.zeros((nworld,), dtype=int)
d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
_allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
return d
@@ -1176,16 +1321,7 @@ def put_data(
# create efc
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:
continue
shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
val = np.zeros(shape, dtype=f.type.dtype)
if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
efc = types.Constraint(**efc_kwargs)
efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS, mjd)
if is_sparse(mjm):
J_rownnz = np.zeros(njmax, dtype=np.int32)
@@ -1236,12 +1372,30 @@ def put_data(
"njmax_nnz": njmax_nnz,
# fields set after initialization:
"solver_niter": None,
"qM": None,
"M": None,
"qLD": None,
"nacon": None,
# island arrays
"nisland": None,
"tree_island": None,
"dof_island": None,
"island_dofadr": None,
"island_nv": None,
"island_nefc": None,
"island_ne": None,
"island_nf": None,
"island_efcadr": None,
"nidof": None,
"map_dof2idof": None,
"map_idof2dof": None,
"map_efc2iefc": None,
"map_iefc2efc": None,
"dof_islandid": None,
"efc_islandid": None,
"iqacc": None,
"iqacc_smooth": None,
"iqfrc_smooth": None,
"iqfrc_constraint": None,
}
for f in dataclasses.fields(types.Data):
if f.name in d_kwargs:
@@ -1256,20 +1410,25 @@ def put_data(
d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
if is_sparse(mjm):
d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float)
if check_version("mujoco>=3.8.1.dev910242375"):
d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.M), dtype=float)
else:
d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.qM[mjm.mapM2M]), dtype=float)
d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float)
else:
qM = np.zeros((mjm.nv, mjm.nv))
mujoco.mj_fullM(mjm, qM, mjd.qM)
qLD = np.linalg.cholesky(qM) if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
M = np.zeros((mjm.nv, mjm.nv))
if check_version("mujoco>=3.8.1.dev910242375"):
mujoco.mju_sym2dense(M, mjd.M, mjm.M_rownnz, mjm.M_rowadr, mjm.M_colind)
qLD = np.linalg.cholesky(M).T if (mjd.M != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
else:
mujoco.mj_fullM(mjm, M, mjd.qM)
qLD = np.linalg.cholesky(M).T if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
padding = sizes["nv_pad"] - mjm.nv
qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
M_padded = np.pad(M, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
d.M = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), M_padded), dtype=float)
d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), 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)
_allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
@@ -1403,6 +1562,9 @@ def get_data_into(
result.qfrc_constraint[:] = d.qfrc_constraint.numpy()[world_id]
result.qfrc_inverse[:] = d.qfrc_inverse.numpy()[world_id]
if mjm.nhistory > 0:
result.history[:] = d.history.numpy()[world_id]
# contact
result.contact.dist[:ncon] = d.contact.dist.numpy()[ncon_filter]
result.contact.pos[:ncon] = d.contact.pos.numpy()[ncon_filter]
@@ -1417,17 +1579,27 @@ def get_data_into(
result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon]
if is_sparse(mjm):
result.qM[:] = d.qM.numpy()[world_id, 0]
if check_version("mujoco>=3.8.1.dev910242375"):
result.M[:] = d.M.numpy()[world_id, 0]
else:
result.qM[mjm.mapM2M] = d.M.numpy()[world_id, 0]
result.qLD[:] = d.qLD.numpy()[world_id, 0]
else:
qM = d.qM.numpy()[world_id]
adr = 0
for i in range(mjm.nv):
j = i
while j >= 0:
result.qM[adr] = qM[i, j]
j = mjm.dof_parentid[j]
adr += 1
M = d.M.numpy()[world_id]
if check_version("mujoco>=3.8.1.dev910242375"):
for i in range(mjm.nv):
adr = mjm.M_rowadr[i]
for k in range(mjm.M_rownnz[i]):
col = mjm.M_colind[adr + k]
result.M[adr + k] = M[i, col]
else:
adr = 0
for i in range(mjm.nv):
j = i
while j >= 0:
result.qM[adr] = M[i, j]
j = mjm.dof_parentid[j]
adr += 1
mujoco.mj_factorM(mjm, result)
if nefc > 0:
@@ -1466,6 +1638,7 @@ def get_data_into(
result.efc_frictionloss[:] = d.efc.frictionloss.numpy()[world_id, efc_idx]
result.efc_state[:] = d.efc.state.numpy()[world_id, efc_idx]
result.efc_force[:] = d.efc.force.numpy()[world_id, efc_idx]
result.efc_island[:] = d.efc.island.numpy()[world_id, efc_idx]
# rne_postconstraint
result.cacc[:] = d.cacc.numpy()[world_id]
@@ -1487,8 +1660,51 @@ def get_data_into(
# islands
nisland = d.nisland.numpy()[world_id]
result.nisland = nisland
if nisland:
if d.tree_island.shape[1] > 0 and nisland:
result.tree_island[:] = d.tree_island.numpy()[world_id]
result.dof_island[:] = d.dof_island.numpy()[world_id]
result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
result.island_ne[:nisland] = d.island_ne.numpy()[world_id, :nisland]
result.island_nf[:nisland] = d.island_nf.numpy()[world_id, :nisland]
result.island_iefcadr[:nisland] = d.island_efcadr.numpy()[world_id, :nisland]
nv = mjm.nv
result.map_dof2idof[:nv] = d.map_dof2idof.numpy()[world_id, :nv]
result.map_idof2dof[:nv] = d.map_idof2dof.numpy()[world_id, :nv]
result.map_efc2iefc[:nefc] = d.map_efc2iefc.numpy()[world_id, :nefc]
result.map_iefc2efc[:nefc] = d.map_iefc2efc.numpy()[world_id, :nefc]
result.iefc_type[:nefc] = d.efc.itype.numpy()[world_id, :nefc]
result.iefc_id[:nefc] = d.efc.iid.numpy()[world_id, :nefc]
result.iefc_D[:nefc] = d.efc.iD.numpy()[world_id, :nefc]
result.iefc_aref[:nefc] = d.efc.iaref.numpy()[world_id, :nefc]
result.iefc_frictionloss[:nefc] = d.efc.ifrictionloss.numpy()[world_id, :nefc]
result.iefc_state[:nefc] = d.efc.istate.numpy()[world_id, :nefc]
result.iefc_force[:nefc] = d.efc.iforce.numpy()[world_id, :nefc]
if is_sparse(mjm):
iefc_J = np.zeros((nefc, mjm.nv))
mujoco.mju_sparse2dense(
iefc_J,
d.efc.iJ.numpy()[world_id, 0],
d.efc.iJ_rownnz.numpy()[world_id, :nefc],
d.efc.iJ_rowadr.numpy()[world_id, :nefc],
d.efc.iJ_colind.numpy()[world_id, 0],
)
else:
iefc_J = d.efc.iJ.numpy()[world_id, :nefc, : mjm.nv]
if mujoco.mj_isSparse(mjm):
mujoco.mju_dense2sparse(
result.iefc_J,
iefc_J,
result.iefc_J_rownnz,
result.iefc_J_rowadr,
result.iefc_J_colind,
)
else:
result.iefc_J[: nefc * mjm.nv] = iefc_J.flatten()
def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@@ -1511,14 +1727,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
xfrc_applied_out[worldid, bodyid][elemid] = 0.0
@wp.kernel(module="unique", enable_backward=False)
def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
def reset_M(reset_in: wp.array[bool], M_out: wp.array3d[float]):
worldid, elemid1, elemid2 = wp.tid()
if wp.static(reset is not None):
if not reset_in[worldid]:
return
qM_out[worldid, elemid1, elemid2] = 0.0
M_out[worldid, elemid1, elemid2] = 0.0
@wp.kernel(module="unique", enable_backward=False)
def reset_nworld(
@@ -1670,10 +1886,10 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
wp.launch(
reset_qM,
dim=(d.nworld, d.qM.shape[1], d.qM.shape[2]),
reset_M,
dim=(d.nworld, d.M.shape[1], d.M.shape[2]),
inputs=[reset_input],
outputs=[d.qM],
outputs=[d.M],
)
# set mocap_pos/quat = body_pos/quat for mocap bodies
@@ -1786,11 +2002,12 @@ def _copy_tendon_length0(
def _compute_meaninertia(
nv: int,
is_sparse: bool,
dof_Madr_in: wp.array[int],
qM_in: wp.array3d[float],
M_rownnz_in: wp.array[int],
M_rowadr_in: wp.array[int],
M_in: wp.array3d[float],
meaninertia_out: wp.array[float],
):
"""Compute mean diagonal inertia from qM at qpos0."""
"""Compute mean diagonal inertia from M at qpos0."""
worldid = wp.tid()
if nv == 0:
@@ -1800,12 +2017,12 @@ def _compute_meaninertia(
total = float(0.0)
for i in range(nv):
if is_sparse:
# Sparse: qM is flattened lower triangular, diagonal at dof_Madr[i]
madr = dof_Madr_in[i]
total += qM_in[worldid, 0, madr]
# Sparse: M is in CSR format, diagonal at M_rowadr_in[i] + M_rownnz_in[i] - 1
madr = M_rowadr_in[i] + M_rownnz_in[i] - 1
total += M_in[worldid, 0, madr]
else:
# Dense: qM is 2D matrix, diagonal at [i,i]
total += qM_in[worldid, i, i]
# Dense: M is 2D matrix, diagonal at [i,i]
total += M_in[worldid, i, i]
meaninertia_out[worldid % meaninertia_out.shape[0]] = total / float(nv)
@@ -2288,11 +2505,11 @@ def set_const_0(m: types.Model, d: types.Data):
smooth.factor_m(m, d)
smooth.transmission(m, d)
# Compute meaninertia from qM diagonal at qpos0
# Compute meaninertia from M diagonal at qpos0
wp.launch(
_compute_meaninertia,
dim=d.nworld,
inputs=[m.nv, m.is_sparse, m.dof_Madr, d.qM],
inputs=[m.nv, m.is_sparse, m.M_rownnz, m.M_rowadr, d.M],
outputs=[m.stat.meaninertia],
)
@@ -2452,27 +2669,31 @@ def set_const(m: types.Model, d: types.Data):
Model fields that can be modified safely with set_const:
Field | Notes
---------------------------------|----------------------------------------------
qpos0, qpos_spring |
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.
dof_armature |
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.
================================== ==============================================
Field Notes
================================== ==============================================
qpos0, qpos_spring
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.
dof_armature
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.
================================== ==============================================
For selective updates, use the sub-functions directly based on what changed:
Modified Field | Call
----------------|------------------
body_mass | set_const
body_gravcomp | set_const_fixed
body_inertia | set_const_0
qpos0 | set_const_0
============== ===============
Modified Field Call
============== ===============
body_mass set_const
body_gravcomp set_const_fixed
body_inertia set_const_0
qpos0 set_const_0
============== ===============
Computes:
- Fixed quantities (via set_const_fixed):
@@ -2621,6 +2842,9 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
else:
val = typ(val)
if attr == "disableflags" and isinstance(obj, types.Option) and not ENABLE_ISLANDS:
val = int(val) | types.DisableBit.ISLAND
setattr(obj, attr, val)
@@ -2724,11 +2948,13 @@ def create_render_context(
render_seg: list[bool] | bool | None = None,
use_textures: bool = True,
use_shadows: bool = False,
use_ambient_lighting: bool = True,
enabled_geom_groups: list[int] = [0, 1, 2],
cam_active: list[bool] | None = None,
flex_render_smooth: bool = True,
use_precomputed_rays: bool = True,
render_skybox: bool = False,
enable_backface_culling: bool = True,
) -> types.RenderContext:
"""Creates a render context on device.
@@ -2743,6 +2969,8 @@ def create_render_context(
If None, uses the MuJoCo model values.
use_textures: Whether to use textures.
use_shadows: Whether to use shadows.
use_ambient_lighting: Whether to add the renderer's hemispheric ambient
lighting term before applying model lights.
enabled_geom_groups: The geom groups to render.
cam_active: List of booleans indicating which cameras to include in rendering.
If None, all cameras are included.
@@ -2751,6 +2979,10 @@ def create_render_context(
When using domain randomization for camera intrinsics, set to False.
render_skybox: Whether to shade missed rays with the MuJoCo skybox texture.
Requires the model to contain a texture with type `mjTEXTURE_SKYBOX`.
enable_backface_culling: Drop primitive-ray hits whose normal faces away from
the ray (ray origin inside the geom). Matches MuJoCo's
mesh-ray rule. Default True. Disable for a small
performance gain when no camera is ever inside a geom.
Returns:
The render context containing rendering fields and output arrays on device.
@@ -2758,13 +2990,7 @@ def create_render_context(
mjd = mujoco.MjData(mjm)
mujoco.mj_forward(mjm, mjd)
constructor = "sah"
if check_version("warp>=1.13.0.dev20260325"):
# TODO: The cubql constructor and is_cubql_available exist only in
# recent Warp 1.13+ builds, modify this after warp is updated to 1.13+.
_cubql_avail = getattr(wp, "is_cubql_available", None)
if callable(_cubql_avail) and _cubql_avail():
constructor = "cubql"
constructor = "cubql"
# Mesh BVHs – build for all meshes so per-world variants are available
nmesh = mjm.nmesh
@@ -2937,6 +3163,7 @@ def create_render_context(
cam_id_map=wp.array(active_cam_indices, dtype=int),
use_textures=use_textures,
use_shadows=use_shadows,
use_ambient_lighting=use_ambient_lighting,
background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
use_precomputed_rays=use_precomputed_rays,
render_skybox=render_skybox,
@@ -2982,6 +3209,7 @@ def create_render_context(
render_seg=wp.array(render_seg, dtype=bool),
znear=znear,
total_rays=int(total),
enable_backface_culling=enable_backface_culling,
)
bvh.build_scene_bvh(mjm, mjd, rc, nworld)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+821 -20
View File
@@ -18,6 +18,7 @@ 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 IslandSolverContext
from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -35,11 +36,15 @@ def _tree_edges(
eq_obj1id: wp.array[int],
eq_obj2id: wp.array[int],
eq_objtype: wp.array[int],
is_sparse: bool,
# Data in:
nefc_in: wp.array[int],
contact_geom_in: wp.array[wp.vec2i],
efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d[int],
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
efc_J_in: wp.array3d[float],
njmax_in: int,
# Out:
@@ -131,26 +136,39 @@ def _tree_edges(
first_tree = int(-1)
has_cross_edge = int(0)
for dof in range(nv):
# TODO(team): sparse efc_J
# TODO(team): tree dof skip
J_val = efc_J_in[worldid, efcid, dof]
if J_val != 0.0:
tree = dof_treeid[dof]
if tree < 0:
count = nv
rowadr = 0
if is_sparse:
count = efc_J_rownnz_in[worldid, efcid]
rowadr = efc_J_rowadr_in[worldid, efcid]
for i in range(count):
dof = i
if is_sparse:
sparseid = rowadr + i
dof = efc_J_colind_in[worldid, 0, sparseid]
else:
J_val = efc_J_in[worldid, efcid, dof]
if J_val == 0.0:
continue
if first_tree == -1:
first_tree = tree
elif tree != first_tree:
t1 = wp.min(first_tree, tree)
t2 = wp.max(first_tree, tree)
wp.atomic_max(tree_tree, worldid, t1, t2, 1)
has_cross_edge = 1
tree = dof_treeid[dof]
if tree < 0:
continue
if first_tree == -1:
first_tree = tree
elif tree != first_tree:
t1 = wp.min(first_tree, tree)
t2 = wp.max(first_tree, tree)
wp.atomic_max(tree_tree, worldid, t1, t2, 1)
wp.atomic_max(tree_tree, worldid, t2, t1, 1)
has_cross_edge = 1
if first_tree >= 0 and has_cross_edge == 0:
wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
@event_scope
def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
"""Compute tree-tree adjacency matrix."""
tree_tree.zero_()
@@ -168,10 +186,14 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
m.eq_obj1id,
m.eq_obj2id,
m.eq_objtype,
m.is_sparse,
d.nefc,
d.contact.geom,
d.efc.type,
d.efc.id,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.efc.J_colind,
d.efc.J,
d.njmax,
],
@@ -242,6 +264,19 @@ def _flood_fill(
nisland_out[worldid] = nisland
@event_scope
def flood_fill(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
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],
)
@event_scope
def island(m: types.Model, d: types.Data):
"""Discover constraint islands."""
@@ -254,12 +289,778 @@ def island(m: types.Model, d: types.Data):
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)
flood_fill(m, d, tree_tree)
@wp.kernel
def _island_count_dofs(
dof_treeid: wp.array[int],
tree_island_in: wp.array2d[int],
dof_island_out: wp.array2d[int],
island_nv_out: wp.array2d[int],
):
worldid, dofid = wp.tid()
island_id = tree_island_in[worldid, dof_treeid[dofid]]
dof_island_out[worldid, dofid] = island_id
if island_id >= 0:
wp.atomic_add(island_nv_out, worldid, island_id, 1)
@wp.kernel
def _island_scan_sizes(
nisland_in: wp.array[int],
island_idofadr_out: wp.array2d[int],
island_nv_inout: wp.array2d[int],
island_nefc_inout: wp.array2d[int],
island_iefcadr_out: wp.array2d[int],
nidof_out: wp.array[int],
):
worldid = wp.tid()
nisland = nisland_in[worldid]
if nisland == 0:
nidof_out[worldid] = 0
return
# Scan DOFs and Constraints
island_idofadr_out[worldid, 0] = 0
island_iefcadr_out[worldid, 0] = 0
for i in range(1, nisland):
island_idofadr_out[worldid, i] = island_idofadr_out[worldid, i - 1] + island_nv_inout[worldid, i - 1]
island_iefcadr_out[worldid, i] = island_iefcadr_out[worldid, i - 1] + island_nefc_inout[worldid, i - 1]
nidof = island_idofadr_out[worldid, nisland - 1] + island_nv_inout[worldid, nisland - 1]
nidof_out[worldid] = nidof
# Reset for recount
for i in range(nisland):
island_nv_inout[worldid, i] = 0
island_nefc_inout[worldid, i] = 0
@wp.kernel
def _island_map_dofs(
nv: int,
dof_island_in: wp.array2d[int],
island_idofadr_in: wp.array2d[int],
nidof_in: wp.array[int],
island_nv_inout: wp.array2d[int],
map_dof2idof_out: wp.array2d[int],
map_idof2dof_out: wp.array2d[int],
idof_islandid_out: wp.array2d[int],
unconstrained_cnt_inout: wp.array2d[int],
):
worldid, dofid = wp.tid()
nidof = nidof_in[worldid]
island_id = dof_island_in[worldid, dofid]
if island_id >= 0:
local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
idof = island_idofadr_in[worldid, island_id] + local_idx
idof_islandid_out[worldid, idof] = island_id
else:
cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
idof = nidof + cnt
map_dof2idof_out[worldid, dofid] = idof
map_idof2dof_out[worldid, idof] = dofid
@wp.kernel
def _island_count_constraints(
nefc_in: wp.array[int],
njmax_in: int,
efc_tree_in: wp.array2d[int],
tree_island_in: wp.array2d[int],
efc_type_in: wp.array2d[int],
efc_island_out: wp.array2d[int],
island_nefc_out: wp.array2d[int],
island_ne_out: wp.array2d[int],
island_nf_out: wp.array2d[int],
):
worldid, efcid = wp.tid()
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
efc_tree = efc_tree_in[worldid, efcid]
if efc_tree < 0:
efc_island_out[worldid, efcid] = -1
return
island_id = tree_island_in[worldid, efc_tree]
efc_island_out[worldid, efcid] = island_id
if island_id >= 0:
wp.atomic_add(island_nefc_out, worldid, island_id, 1)
efc_type = efc_type_in[worldid, efcid]
if efc_type == ConstraintType.EQUALITY:
wp.atomic_add(island_ne_out, worldid, island_id, 1)
elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
wp.atomic_add(island_nf_out, worldid, island_id, 1)
@wp.kernel
def _island_map_constraints(
nefc_in: wp.array[int],
njmax_in: int,
efc_island_in: wp.array2d[int],
island_iefcadr_in: wp.array2d[int],
island_ne_in: wp.array2d[int],
island_nf_in: wp.array2d[int],
efc_type_in: wp.array2d[int],
# Counters (inout):
island_ne_mapped_inout: wp.array2d[int],
island_nf_mapped_inout: wp.array2d[int],
island_nother_mapped_inout: wp.array2d[int],
island_nefc_inout: wp.array2d[int],
# Out:
map_efc2iefc_out: wp.array2d[int],
map_iefc2efc_out: wp.array2d[int],
iefc_islandid_out: wp.array2d[int],
):
worldid, efcid = wp.tid()
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
island_id = efc_island_in[worldid, efcid]
if island_id >= 0:
efc_type = efc_type_in[worldid, efcid]
# 1. Determine local index and absolute index ic based on category
if efc_type == ConstraintType.EQUALITY:
local_idx = wp.atomic_add(island_ne_mapped_inout, worldid, island_id, 1)
ic = island_iefcadr_in[worldid, island_id] + local_idx
elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
local_idx = wp.atomic_add(island_nf_mapped_inout, worldid, island_id, 1)
ic = island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + local_idx
else:
local_idx = wp.atomic_add(island_nother_mapped_inout, worldid, island_id, 1)
ic = (
island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + island_nf_in[worldid, island_id] + local_idx
)
# 2. Increment overall island_nefc counter to reconstruct d.island_nefc
wp.atomic_add(island_nefc_inout, worldid, island_id, 1)
# 3. Store mappings
map_efc2iefc_out[worldid, efcid] = ic
map_iefc2efc_out[worldid, ic] = efcid
iefc_islandid_out[worldid, ic] = island_id
@wp.kernel
def _island_scan_sparse_rows(
nisland_in: wp.array[int],
efc_J_rownnz_in: wp.array2d[int],
island_nefc_in: wp.array2d[int],
island_iefcadr_in: wp.array2d[int],
map_iefc2efc_in: wp.array2d[int],
iefc_J_rownnz_out: wp.array2d[int],
iefc_J_rowadr_out: wp.array2d[int],
):
worldid = wp.tid()
nisland = nisland_in[worldid]
if nisland == 0:
return
total_gathered_efc = island_iefcadr_in[worldid, nisland - 1] + island_nefc_in[worldid, nisland - 1]
running_rowadr = int(0)
for ic in range(total_gathered_efc):
c = map_iefc2efc_in[worldid, ic]
rownnz = efc_J_rownnz_in[worldid, c]
iefc_J_rowadr_out[worldid, ic] = running_rowadr
iefc_J_rownnz_out[worldid, ic] = rownnz
running_rowadr += rownnz
@wp.kernel
def _compute_efc_tree(
# Model:
nv: int,
body_treeid: wp.array[int],
jnt_dofadr: wp.array[int],
dof_treeid: wp.array[int],
geom_bodyid: wp.array[int],
site_bodyid: wp.array[int],
eq_type: wp.array[int],
eq_obj1id: wp.array[int],
eq_obj2id: wp.array[int],
eq_objtype: wp.array[int],
is_sparse: bool,
# Data in:
nefc_in: wp.array[int],
contact_geom_in: wp.array[wp.vec2i],
efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d[int],
efc_J_in: wp.array3d[float],
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
njmax_in: int,
# Out:
efc_tree_out: wp.array2d[int],
):
"""Compute the first non-negative tree for each constraint."""
worldid, efcid = wp.tid()
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
efc_type = efc_type_in[worldid, efcid]
efc_id = efc_id_in[worldid, efcid]
tree = int(-1)
use_generic = int(0)
# equality (connect/weld)
if efc_type == ConstraintType.EQUALITY:
eq_t = eq_type[efc_id]
if eq_t == EqType.CONNECT or eq_t == EqType.WELD:
b1 = eq_obj1id[efc_id]
b2 = eq_obj2id[efc_id]
if eq_objtype[efc_id] == ObjType.SITE:
b1 = site_bodyid[b1]
b2 = site_bodyid[b2]
t1 = body_treeid[b1]
t2 = body_treeid[b2]
if t1 >= 0:
tree = t1
else:
tree = t2
else:
# JOINT, TENDON, FLEX: generic scan
use_generic = 1
# joint friction
elif efc_type == ConstraintType.FRICTION_DOF:
tree = dof_treeid[efc_id]
# joint limit
elif efc_type == ConstraintType.LIMIT_JOINT:
tree = dof_treeid[jnt_dofadr[efc_id]]
# contact
elif (
efc_type == ConstraintType.CONTACT_FRICTIONLESS
or efc_type == ConstraintType.CONTACT_PYRAMIDAL
or efc_type == ConstraintType.CONTACT_ELLIPTIC
):
geom_pair = contact_geom_in[efc_id]
g1 = geom_pair[0]
g2 = geom_pair[1]
if g1 >= 0 and g2 >= 0:
t1 = body_treeid[geom_bodyid[g1]]
t2 = body_treeid[geom_bodyid[g2]]
if t1 >= 0:
tree = t1
else:
tree = t2
else:
# flex contacts: generic scan
use_generic = 1
else:
# generic: scan Jacobian row
use_generic = 1
if use_generic:
count = nv
rowadr = 0
if is_sparse:
count = efc_J_rownnz_in[worldid, efcid]
rowadr = efc_J_rowadr_in[worldid, efcid]
for i in range(count):
dof = i
if is_sparse:
sparseid = rowadr + i
dof = efc_J_colind_in[worldid, 0, sparseid]
else:
J_val = efc_J_in[worldid, efcid, dof]
if J_val == 0.0:
continue
t = dof_treeid[dof]
if t >= 0:
tree = t
break
efc_tree_out[worldid, efcid] = tree
@wp.kernel
def _gather_efc_and_jacobian(
# Model:
is_sparse: bool,
# Data in:
nefc_in: wp.array[int],
efc_D_in: wp.array2d[float],
efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d[int],
efc_frictionloss_in: wp.array2d[float],
efc_aref_in: wp.array2d[float],
efc_J_in: wp.array3d[float],
# Sparse Jacobian arrays:
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
iefc_J_rowadr_in: wp.array2d[int],
# In:
njmax_in: int,
map_iefc2efc_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
map_dof2idof_in: wp.array2d[int],
nidof_in: wp.array[int],
# Out:
iefc_D_out: wp.array2d[float],
iefc_type_out: wp.array2d[int],
iefc_id_out: wp.array2d[int],
iefc_frictionloss_out: wp.array2d[float],
iefc_aref_out: wp.array2d[float],
iefc_J_out: wp.array3d[float],
iefc_J_colind_out: wp.array3d[int],
):
"""Gather constraint arrays and Jacobian into island-local dense order."""
worldid, iefcid = wp.tid()
if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
return
c = map_iefc2efc_in[worldid, iefcid]
# Scalar constraint fields
iefc_D_out[worldid, iefcid] = efc_D_in[worldid, c]
iefc_type_out[worldid, iefcid] = efc_type_in[worldid, c]
iefc_id_out[worldid, iefcid] = efc_id_in[worldid, c]
iefc_frictionloss_out[worldid, iefcid] = efc_frictionloss_in[worldid, c]
iefc_aref_out[worldid, iefcid] = efc_aref_in[worldid, c]
# Jacobian: convert to island-local dense format
nid = nidof_in[worldid]
if is_sparse:
rownnz = efc_J_rownnz_in[worldid, c]
rowadr_in = efc_J_rowadr_in[worldid, c]
rowadr_out = iefc_J_rowadr_in[worldid, iefcid]
for i in range(rownnz):
sparseid_in = rowadr_in + i
sparseid_out = rowadr_out + i
dof = efc_J_colind_in[worldid, 0, sparseid_in]
idof = map_dof2idof_in[worldid, dof]
# Store in sparse iefc_J_out and iefc_J_colind_out
iefc_J_out[worldid, 0, sparseid_out] = efc_J_in[worldid, 0, sparseid_in]
iefc_J_colind_out[worldid, 0, sparseid_out] = idof
else:
# Dense path: reorder rows by iefc, columns by idof
for idof in range(nid):
dof = map_idof2dof_in[worldid, idof]
iefc_J_out[worldid, iefcid, idof] = efc_J_in[worldid, c, dof]
@wp.kernel
def _gather_dof_arrays(
# Data in:
qacc_in: wp.array2d[float],
qacc_smooth_in: wp.array2d[float],
qfrc_smooth_in: wp.array2d[float],
# In:
nidof_in: wp.array[int],
map_idof2dof_in: wp.array2d[int],
# Out:
iacc_out: wp.array2d[float],
iacc_smooth_out: wp.array2d[float],
ifrc_smooth_out: wp.array2d[float],
):
"""Gather DOF arrays into island-local order."""
worldid, idofid = wp.tid()
if idofid >= nidof_in[worldid]:
return
dof = map_idof2dof_in[worldid, idofid]
iacc_out[worldid, idofid] = qacc_in[worldid, dof]
iacc_smooth_out[worldid, idofid] = qacc_smooth_in[worldid, dof]
ifrc_smooth_out[worldid, idofid] = qfrc_smooth_in[worldid, dof]
@wp.kernel
def _scatter_dof_arrays(
# In:
qacc_smooth_in: wp.array2d[float],
qfrc_smooth_in: wp.array2d[float],
dof_island_in: wp.array2d[int],
iacc_in: wp.array2d[float],
ifrc_constraint_in: wp.array2d[float],
iMa_in: wp.array2d[float],
map_dof2idof_in: wp.array2d[int],
scatter_Ma: bool,
# Out:
qacc_out: wp.array2d[float],
qfrc_constraint_out: wp.array2d[float],
Ma_out: wp.array2d[float],
):
"""Scatter island results to global arrays, copy qacc_smooth for non-island DOFs."""
worldid, dofid = wp.tid()
if dof_island_in[worldid, dofid] < 0:
qacc_out[worldid, dofid] = qacc_smooth_in[worldid, dofid]
qfrc_constraint_out[worldid, dofid] = 0.0
if scatter_Ma:
Ma_out[worldid, dofid] = qfrc_smooth_in[worldid, dofid]
else:
idof = map_dof2idof_in[worldid, dofid]
qacc_out[worldid, dofid] = iacc_in[worldid, idof]
qfrc_constraint_out[worldid, dofid] = ifrc_constraint_in[worldid, idof]
if scatter_Ma:
Ma_out[worldid, dofid] = iMa_in[worldid, idof]
@wp.kernel
def _scatter_efc_arrays(
# In:
nefc_in: wp.array[int],
njmax_in: int,
map_iefc2efc_in: wp.array2d[int],
iefc_force_in: wp.array2d[float],
iefc_state_in: wp.array2d[int],
# Out:
efc_force_out: wp.array2d[float],
efc_state_out: wp.array2d[int],
):
"""Scatter island-local constraint results back to global order."""
worldid, iefcid = wp.tid()
if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
return
c = map_iefc2efc_in[worldid, iefcid]
efc_force_out[worldid, c] = iefc_force_in[worldid, iefcid]
efc_state_out[worldid, c] = iefc_state_in[worldid, iefcid]
@wp.kernel
def _init_island_arrays(
island_idofadr_out: wp.array2d[int],
island_nv_out: wp.array2d[int],
island_nefc_out: wp.array2d[int],
island_ne_out: wp.array2d[int],
island_nf_out: wp.array2d[int],
island_iefcadr_out: wp.array2d[int],
nidof_out: wp.array[int],
):
worldid, islandid = wp.tid()
island_nv_out[worldid, islandid] = 0
island_nefc_out[worldid, islandid] = 0
island_ne_out[worldid, islandid] = 0
island_nf_out[worldid, islandid] = 0
island_idofadr_out[worldid, islandid] = 0
island_iefcadr_out[worldid, islandid] = 0
if islandid == 0:
nidof_out[worldid] = 0
@wp.kernel
def _init_dof_arrays(
dof_island_out: wp.array2d[int],
map_dof2idof_out: wp.array2d[int],
map_idof2dof_out: wp.array2d[int],
idof_islandid_out: wp.array2d[int],
):
worldid, dofid = wp.tid()
dof_island_out[worldid, dofid] = -1
map_dof2idof_out[worldid, dofid] = 0
map_idof2dof_out[worldid, dofid] = 0
idof_islandid_out[worldid, dofid] = -1
@wp.kernel
def _init_efc_arrays(
efc_island_out: wp.array2d[int],
map_efc2iefc_out: wp.array2d[int],
map_iefc2efc_out: wp.array2d[int],
iefc_islandid_out: wp.array2d[int],
efc_tree_out: wp.array2d[int],
):
worldid, efcid = wp.tid()
efc_island_out[worldid, efcid] = -1
map_efc2iefc_out[worldid, efcid] = 0
map_iefc2efc_out[worldid, efcid] = 0
iefc_islandid_out[worldid, efcid] = -1
efc_tree_out[worldid, efcid] = -1
@event_scope
def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverContext):
"""Compute DOF/constraint island mappings after island discovery.
Populates island solver context arrays via ctx: nv, nefc, ne, nf,
iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
and d.island_dofadr.
Args:
m: Model.
d: Data.
ctx: IslandSolverContext.
"""
# Ensure dof_islandid / efc_islandid are allocated at the right shape
if d.dof_islandid.shape[1] != m.nv:
d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
if d.efc_islandid.shape[1] != d.njmax:
d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
# Ensure island-local DOF arrays are allocated at the right shape
if d.iqacc.shape[1] != m.nv:
nw = d.nworld
d.iqacc = wp.empty((nw, m.nv), dtype=float)
d.iqacc_smooth = wp.empty((nw, m.nv), dtype=float)
d.iqfrc_smooth = wp.empty((nw, m.nv), dtype=float)
d.iqfrc_constraint = wp.empty((nw, m.nv), dtype=float)
wp.launch(
_init_island_arrays,
dim=(d.nworld, m.ntree),
inputs=[],
outputs=[
d.island_dofadr,
d.island_nv,
d.island_nefc,
d.island_ne,
d.island_nf,
d.island_efcadr,
d.nidof,
],
)
wp.launch(
_init_dof_arrays,
dim=(d.nworld, m.nv),
inputs=[],
outputs=[d.dof_island, d.map_dof2idof, d.map_idof2dof, d.dof_islandid],
)
efc_tree = wp.empty((d.nworld, d.njmax), dtype=int)
wp.launch(
_init_efc_arrays,
dim=(d.nworld, d.njmax),
inputs=[],
outputs=[d.efc.island, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid, efc_tree],
)
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],
_compute_efc_tree,
dim=(d.nworld, d.njmax),
inputs=[
m.nv,
m.body_treeid,
m.jnt_dofadr,
m.dof_treeid,
m.geom_bodyid,
m.site_bodyid,
m.eq_type,
m.eq_obj1id,
m.eq_obj2id,
m.eq_objtype,
m.is_sparse,
d.nefc,
d.contact.geom,
d.efc.type,
d.efc.id,
d.efc.J,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.efc.J_colind,
d.njmax,
],
outputs=[efc_tree],
)
# 1. Count DOFs per island
wp.launch(
_island_count_dofs,
dim=(d.nworld, m.nv),
inputs=[m.dof_treeid, d.tree_island],
outputs=[d.dof_island, d.island_nv],
)
# 2. Count Constraints per island
wp.launch(
_island_count_constraints,
dim=(d.nworld, d.njmax),
inputs=[d.nefc, d.njmax, efc_tree, d.tree_island, d.efc.type],
outputs=[d.efc.island, d.island_nefc, d.island_ne, d.island_nf],
)
# 3. Scan sizes and reset counters for mapping
wp.launch(
_island_scan_sizes,
dim=d.nworld,
inputs=[d.nisland],
outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
)
# 4. Map DOFs
unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
wp.launch(
_island_map_dofs,
dim=(d.nworld, m.nv),
inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
)
# 5. Map Constraints
ne_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
nf_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
nother_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
wp.launch(
_island_map_constraints,
dim=(d.nworld, d.njmax),
inputs=[
d.nefc,
d.njmax,
d.efc.island,
d.island_efcadr,
d.island_ne,
d.island_nf,
d.efc.type,
ne_mapped,
nf_mapped,
nother_mapped,
],
outputs=[d.island_nefc, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid],
)
# 6. Scan Sparse Rows (if sparse)
if m.is_sparse:
wp.launch(
_island_scan_sparse_rows,
dim=d.nworld,
inputs=[d.nisland, d.efc.J_rownnz, d.island_nefc, d.island_efcadr, d.map_iefc2efc],
outputs=[d.efc.iJ_rownnz, d.efc.iJ_rowadr],
)
@event_scope
def gather_island_inputs(m: types.Model, d: types.Data, ctx: IslandSolverContext):
"""Gather constraint and DOF arrays into island-local order.
Populates d.iefc (D, type, id, frictionloss, aref, J, J_colind) and
d.iqacc, d.iqacc_smooth, d.iqfrc_smooth.
Must be called after compute_island_mapping() and before per-island solving.
Args:
m: Model.
d: Data.
ctx: IslandSolverContext whose arrays are populated.
"""
# Gather constraint arrays and dense Jacobian (fused)
wp.launch(
_gather_efc_and_jacobian,
dim=(d.nworld, d.njmax),
inputs=[
m.is_sparse,
d.nefc,
d.efc.D,
d.efc.type,
d.efc.id,
d.efc.frictionloss,
d.efc.aref,
d.efc.J,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.efc.J_colind,
d.efc.iJ_rowadr,
d.njmax,
d.map_iefc2efc,
d.map_idof2dof,
d.map_dof2idof,
d.nidof,
],
outputs=[
d.efc.iD,
d.efc.itype,
d.efc.iid,
d.efc.ifrictionloss,
d.efc.iaref,
d.efc.iJ,
d.efc.iJ_colind,
],
)
# Gather DOF arrays
wp.launch(
_gather_dof_arrays,
dim=(d.nworld, m.nv),
inputs=[
d.qacc,
d.qacc_smooth,
d.qfrc_smooth,
d.nidof,
d.map_idof2dof,
],
outputs=[
d.iqacc,
d.iqacc_smooth,
d.iqfrc_smooth,
],
)
@event_scope
def scatter_island_results(m: types.Model, d: types.Data, ctx: IslandSolverContext, scatter_Ma: bool):
"""Scatter island-local solver results back to global arrays.
Reads ctx qacc, qfrc_constraint, Ma, d.efc iforce, istate and
writes them back to d.qacc, d.qfrc_constraint, d.efc.Ma, d.efc.force,
d.efc.state. Unconstrained DOFs receive qacc_smooth and zero qfrc.
Args:
m: Model.
d: Data.
ctx: IslandSolverContext which contains results.
scatter_Ma: Whether to scatter Ma for Euler/implicit integrators.
"""
# Scatter DOF results (and optionally Ma for Euler/implicit integrators)
wp.launch(
_scatter_dof_arrays,
dim=(d.nworld, m.nv),
inputs=[
d.qacc_smooth,
d.qfrc_smooth,
d.dof_island,
d.iqacc,
d.iqfrc_constraint,
ctx.Ma,
d.map_dof2idof,
scatter_Ma,
],
outputs=[
d.qacc,
d.qfrc_constraint,
d.efc.Ma,
],
)
# Scatter constraint results (force and state from island-local arrays)
wp.launch(
_scatter_efc_arrays,
dim=(d.nworld, d.njmax),
inputs=[
d.nefc,
d.njmax,
d.map_iefc2efc,
d.efc.iforce,
d.efc.istate,
],
outputs=[
d.efc.force,
d.efc.state,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+13 -1
View File
@@ -707,10 +707,14 @@ def ray_mesh_with_bvh(
pnt: wp.vec3,
vec: wp.vec3,
max_t: float,
cull_backfaces: bool,
) -> Tuple[float, wp.vec3, float, float, int, int]:
"""Returns intersection information for ray mesh intersections.
Requires wp.Mesh be constructed and their ids to be passed.
When ``cull_backfaces`` is True, the function rejects exit-face hits in the
local-space frame. This matches MuJoCo OpenGL rendering's backface culling rule.
"""
t = float(-1.0)
u = float(0.0)
@@ -722,7 +726,12 @@ def ray_mesh_with_bvh(
lpnt, lvec = _ray_map(pos, mat, pnt, vec)
hit = wp.mesh_query_ray(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t, t, u, v, sign, n, f)
if hit and wp.dot(lvec, n) < 0.0: # Backface culling in local space
if not hit:
return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
# Front-face hit, or back-face hit when cull is disabled: rotate the
# local-space normal into world space and return the hit.
if (not cull_backfaces) or wp.dot(lvec, n) < 0.0:
normal = mat @ n
normal = wp.normalize(normal)
return t, normal, u, v, f, mesh_geom_id
@@ -1055,6 +1064,8 @@ def _ray_geom_mesh_bvh(
if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
# Public ray API (mjw.ray / mjw.rays) preserves MuJoCo's mj_ray cull rule:
# rangefinder sensors and user-facing ray casts always cull back-faces.
t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
bvh_ids,
geom_dataid[worldid % geom_dataid.shape[0], geomid],
@@ -1063,6 +1074,7 @@ def _ray_geom_mesh_bvh(
pnt,
vec,
min_dist,
True,
)
if t >= 0.0 and t < min_dist:
return t, n
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+27 -6
View File
@@ -180,6 +180,7 @@ def cast_ray(
flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
cull_backfaces: bool,
) -> Tuple[int, float, wp.vec3, float, float, int, int]:
dist = float(MJ_MAXVAL)
normal = wp.vec3(0.0, 0.0, 0.0)
@@ -232,6 +233,7 @@ def cast_ray(
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
@@ -281,6 +283,7 @@ def cast_ray(
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if gtype == GeomType.FLEX:
hit_geom_id = -2
@@ -308,6 +311,11 @@ def cast_ray(
if d >= 0.0:
hit_mesh_id = flexid
# Backface cull: drop exit-face hits when the ray origin is inside the geom,
# matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule.
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
if d >= 0.0 and d < dist:
dist = d
normal = n
@@ -349,6 +357,7 @@ def cast_ray_first_hit(
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
max_dist: float,
cull_backfaces: bool,
) -> bool:
"""A simpler version of casting rays that only checks for the first hit."""
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
@@ -387,6 +396,7 @@ def cast_ray_first_hit(
ray_origin_world,
ray_dir_world,
max_dist,
cull_backfaces,
)
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
@@ -467,6 +477,11 @@ def cast_ray_first_hit(
)
d = 0.0 if hit else -1.0
# Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent
# hits and skips branches with a zero-vector normal (mesh/flex anyhit).
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
if d >= 0.0 and d < max_dist:
return True
@@ -507,6 +522,7 @@ def compute_lighting(
lightdir: wp.vec3,
normal: wp.vec3,
hitpoint: wp.vec3,
cull_backfaces: bool,
) -> float:
light_contribution = float(0.0)
@@ -571,6 +587,7 @@ def compute_lighting(
shadow_origin,
L,
max_t,
cull_backfaces,
)
if shadow_hit:
@@ -725,6 +742,7 @@ def render(m: Model, d: Data, rc: RenderContext):
flex_group_root,
ray_origin_world,
ray_dir_world,
wp.static(rc.enable_backface_culling),
)
if render_seg[cam_idx] and geom_id != -1:
@@ -798,12 +816,14 @@ def render(m: Model, d: Data, rc: RenderContext):
)
base_color = wp.cw_mul(base_color, tex_color)
len_n = wp.length(normal)
n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
n = wp.normalize(n)
hemispheric = 0.5 * (n[2] + 1.0)
ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
result = 0.5 * wp.cw_mul(base_color, ambient_color)
result = wp.vec3(0.0, 0.0, 0.0)
if wp.static(rc.use_ambient_lighting):
len_n = wp.length(normal)
n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
n = wp.normalize(n)
hemispheric = 0.5 * (n[2] + 1.0)
ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
result = 0.5 * wp.cw_mul(base_color, ambient_color)
# Apply lighting and shadows
for l in range(wp.static(m.nlight)):
@@ -837,6 +857,7 @@ def render(m: Model, d: Data, rc: RenderContext):
light_xdir_in[worldid, l],
normal,
hit_point,
wp.static(rc.enable_backface_culling),
)
result = result + base_color * light_contribution
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+26 -2
View File
@@ -17,6 +17,7 @@ from typing import Any, Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import history
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
@@ -26,6 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import TACTILE_DEPTH_SEMANTICS
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
@@ -917,6 +919,10 @@ def sensor_pos(m: Model, d: Data):
],
)
# apply sensor delay/interval for position sensors
history.apply_sensor_delay(m, d, m.sensor_pos_adr)
history.apply_sensor_delay(m, d, m.sensor_limitpos_adr)
if m.callback.sensor:
m.callback.sensor(m, d, Stage.POS)
@@ -1459,6 +1465,10 @@ def sensor_vel(m: Model, d: Data):
],
)
# apply sensor delay/interval for velocity sensors
history.apply_sensor_delay(m, d, m.sensor_vel_adr)
history.apply_sensor_delay(m, d, m.sensor_limitvel_adr)
if m.callback.sensor:
m.callback.sensor(m, d, Stage.VEL)
@@ -2253,13 +2263,23 @@ def _sensor_tactile(
vel_rel = vel_sensor - vel_other
forceT = wp.vec3(0.0, 0.0, 0.0)
forceT[0] = -depth
if wp.static(TACTILE_DEPTH_SEMANTICS):
forceT[0] = -depth
else:
kMaxDepth = 0.05
pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
force = wp.mul(normal, pressure)
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_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
if wp.static(TACTILE_DEPTH_SEMANTICS):
wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
else:
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])
@@ -2705,6 +2725,10 @@ def sensor_acc(m: Model, d: Data):
],
)
# apply sensor delay/interval for acceleration sensors
history.apply_sensor_delay(m, d, m.sensor_acc_adr)
history.apply_sensor_delay(m, d, m.sensor_limitfrc_adr)
if m.callback.sensor:
m.callback.sensor(m, d, Stage.ACC)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+179 -62
View File
@@ -823,37 +823,38 @@ def _crb_accumulate(
@wp.kernel
def _qM_sparse(
def _M_sparse(
# Model:
dof_bodyid: wp.array[int],
dof_parentid: wp.array[int],
dof_Madr: wp.array[int],
dof_armature: wp.array2d[float],
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
# Data in:
cdof_in: wp.array2d[wp.spatial_vector],
crb_in: wp.array2d[vec10],
# Data out:
qM_out: wp.array3d[float],
M_out: wp.array3d[float],
):
worldid, dofid = wp.tid()
madr_ij = dof_Madr[dofid] # dof_Madr is not batched
bodyid = dof_bodyid[dofid]
madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
# init M(i,i) with armature inertia
qM_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
M_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
# precompute buf = crb_body_i * cdof_i
buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
# sparse backward pass over ancestors
while dofid >= 0:
qM_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
madr_ij += 1
M_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
madr_ij -= 1
dofid = dof_parentid[dofid]
@wp.kernel
def _qM_dense(
def _M_dense(
# Model:
dof_bodyid: wp.array[int],
dof_parentid: wp.array[int],
@@ -862,7 +863,7 @@ def _qM_dense(
cdof_in: wp.array2d[wp.spatial_vector],
crb_in: wp.array2d[vec10],
# Data out:
qM_out: wp.array3d[float],
M_out: wp.array3d[float],
):
worldid, dofid = wp.tid()
bodyid = dof_bodyid[dofid]
@@ -873,15 +874,15 @@ def _qM_dense(
buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
M += wp.dot(cdof_in[worldid, dofid], buf)
qM_out[worldid, dofid, dofid] = M
M_out[worldid, dofid, dofid] = M
# sparse backward pass over ancestors
dofidi = dofid
dofid = dof_parentid[dofid]
while dofid >= 0:
qMij = wp.dot(cdof_in[worldid, dofid], buf)
qM_out[worldid, dofidi, dofid] += qMij
qM_out[worldid, dofid, dofidi] += qMij
Mij = wp.dot(cdof_in[worldid, dofid], buf)
M_out[worldid, dofidi, dofid] += Mij
M_out[worldid, dofid, dofidi] += Mij
dofid = dof_parentid[dofid]
@@ -898,17 +899,17 @@ def crb(m: Model, d: Data):
body_tree = m.body_tree[i]
wp.launch(_crb_accumulate, dim=(d.nworld, body_tree.size), inputs=[m.body_parentid, d.crb, body_tree], outputs=[d.crb])
d.qM.zero_()
d.M.zero_()
if m.is_sparse:
wp.launch(
_qM_sparse,
_M_sparse,
dim=(d.nworld, m.nv),
inputs=[m.dof_bodyid, m.dof_parentid, m.dof_Madr, m.dof_armature, d.cdof, d.crb],
outputs=[d.qM],
inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, m.M_rownnz, m.M_rowadr, d.cdof, d.crb],
outputs=[d.M],
)
else:
wp.launch(
_qM_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.qM]
_M_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.M]
)
@@ -916,16 +917,17 @@ def crb(m: Model, d: Data):
def _tendon_armature(
# Model:
dof_parentid: wp.array[int],
dof_Madr: wp.array[int],
ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array[int],
tendon_armature: wp.array2d[float],
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
is_sparse: bool,
# Data in:
ten_J_in: wp.array2d[float],
# Data out:
qM_out: wp.array3d[float],
M_out: wp.array3d[float],
):
worldid, tenid, dofid = wp.tid()
@@ -947,7 +949,7 @@ def _tendon_armature(
return
if is_sparse:
madr_ij = dof_Madr[dofid]
madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
# sparse backward pass over ancestors
dofidi = dofid
@@ -967,52 +969,40 @@ def _tendon_armature(
else:
ten_Jj = float(0.0)
qMij = armature * ten_Jj * ten_Ji
Mij = armature * ten_Jj * ten_Ji
if is_sparse:
wp.atomic_add(qM_out[worldid, 0], madr_ij, qMij)
madr_ij += 1
wp.atomic_add(M_out[worldid, 0], madr_ij, Mij)
madr_ij -= 1
else:
wp.atomic_add(qM_out[worldid, dofidi], dofid, qMij)
wp.atomic_add(M_out[worldid, dofidi], dofid, Mij)
if dofidi != dofid:
wp.atomic_add(qM_out[worldid, dofid], dofidi, qMij)
wp.atomic_add(M_out[worldid, dofid], dofidi, Mij)
dofid = dof_parentid[dofid]
@event_scope
def tendon_armature(m: Model, d: Data):
"""Add tendon armature to qM."""
"""Add tendon armature to M."""
wp.launch(
_tendon_armature,
dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
inputs=[
m.dof_parentid,
m.dof_Madr,
m.ten_J_rownnz,
m.ten_J_rowadr,
m.ten_J_colind,
m.tendon_armature,
m.M_rownnz,
m.M_rowadr,
m.is_sparse,
d.ten_J,
],
outputs=[d.qM],
outputs=[d.M],
)
@wp.kernel
def _copy_CSR(
# Model:
mapM2M: wp.array[int],
# In:
M_in: wp.array3d[float],
# Out:
L_out: wp.array3d[float],
):
worldid, ind = wp.tid()
L_out[worldid, 0, ind] = M_in[worldid, 0, mapM2M[ind]]
@wp.kernel
def _qLD_acc(
# Model:
@@ -1055,7 +1045,7 @@ def _qLDiag_div(
def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[float], D: wp.array2d[float]):
"""Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
wp.launch(_copy_CSR, dim=(d.nworld, m.nC), inputs=[m.mapM2M, M], outputs=[L])
wp.copy(L, M)
for i in reversed(range(len(m.qLD_updates))):
qLD_updates = m.qLD_updates[i]
@@ -1071,7 +1061,7 @@ def _tile_cholesky_factorize(tile: TileSet):
@wp.kernel(module="unique", enable_backward=False)
def cholesky_factorize(
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
adr: wp.array[int],
# Out:
@@ -1081,8 +1071,8 @@ def _tile_cholesky_factorize(tile: TileSet):
TILE_SIZE = wp.static(tile.size)
dofid = adr[nodeid]
M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
L_tile = wp.tile_cholesky(M_tile)
M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
wp.tile_store(L_out[worldid], L_tile, offset=(dofid, dofid))
return cholesky_factorize
@@ -1090,7 +1080,7 @@ def _tile_cholesky_factorize(tile: TileSet):
def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
"""Dense Cholesky factorization of inertia-like matrix M, assumed spd."""
for tile in m.qM_tiles:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_factorize(tile),
dim=(d.nworld, tile.adr.size),
@@ -1104,9 +1094,9 @@ def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
def factor_m(m: Model, d: Data):
"""Factorization of inertia-like matrix M, assumed spd."""
if m.is_sparse:
_factor_i_sparse(m, d, d.qM, d.qLD, d.qLDiagInv)
_factor_i_sparse(m, d, d.M, d.qLD, d.qLDiagInv)
else:
_factor_i_dense(m, d, d.qM, d.qLD)
_factor_i_dense(m, d, d.M, d.qLD)
@wp.kernel
@@ -2842,15 +2832,15 @@ def _tile_cholesky_solve(tile: TileSet):
dofid = adr[nodeid]
y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
wp.tile_store(x[worldid], x_slice, offset=(dofid,))
return cholesky_solve
def _solve_LD_dense(m: Model, d: Data, L: wp.array3d[float], x: wp.array2d[float], y: wp.array2d[float]):
"""Computes dense backsubstitution: x = inv(L'*L)*y."""
for tile in m.qM_tiles:
"""Computes dense backsubstitution: x = inv(U.T @ U) * y."""
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_solve(tile),
dim=(d.nworld, tile.adr.size),
@@ -2868,9 +2858,9 @@ def solve_LD(
x: wp.array2d[float],
y: wp.array2d[float],
):
"""Computes backsubstitution to solve a linear system of the form x = inv(L'*D*L) * y.
"""Computes backsubstitution for the inertia factorization.
L and D are the factors from the Cholesky factorization of the inertia matrix.
Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
This function dispatches to either a sparse or dense solver depending on Model options.
@@ -2907,8 +2897,9 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
@wp.kernel(module="unique", enable_backward=False)
def cholesky_factorize_solve(
# Data in:
M_in: wp.array3d[float],
# In:
M: wp.array3d[float],
y: wp.array2d[float],
adr: wp.array[int],
# Out:
@@ -2919,12 +2910,12 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
TILE_SIZE = wp.static(tile.size)
dofid = adr[nodeid]
M_tile = wp.tile_load(M[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_cholesky(M_tile)
L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
wp.tile_store(L[worldid], L_tile, offset=(dofid, dofid))
x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
wp.tile_store(x[worldid], x_slice, offset=(dofid,))
return cholesky_factorize_solve
@@ -2938,7 +2929,7 @@ def _factor_solve_i_dense(
y: wp.array2d[float],
L: wp.array3d[float],
):
for tile in m.qM_tiles:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_factorize_solve(tile),
dim=(d.nworld, tile.adr.size),
@@ -2949,9 +2940,9 @@ def _factor_solve_i_dense(
def factor_solve_i(m, d, M, L, D, x, y):
"""Factorizes and solves the linear system: x = inv(L'*D*L) * y or x = inv(L'*L) * y.
"""Factorizes and solves the inertia-like linear system.
M is an inertia-like matrix and L, D are its Cholesky-like factors.
Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
This function first factorizes the matrix M (sparse or dense depending on model options),
then solves the system for x given right-hand side y.
@@ -2960,7 +2951,7 @@ def factor_solve_i(m, d, M, L, D, x, y):
m: The model containing factorization and sparsity information.
d: The data object containing workspace and factorization results.
M: The inertia-like matrix to factorize.
L: Output lower-triangular factor from the factorization (sparse or dense).
L: Output sparse factor or dense upper Cholesky factor.
D: Output diagonal factor from the factorization (only used for sparse).
x: Output array for the solution.
y: Input right-hand side array.
@@ -2972,6 +2963,132 @@ def factor_solve_i(m, d, M, L, D, x, y):
_factor_solve_i_dense(m, d, M, x, y, L)
@cache_kernel
def _factor_solve_lu_sparse_fused(nv: int):
"""Fused sparse LU factorization and solve in a single kernel."""
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
D_rownnz: wp.array[int],
D_rowadr: wp.array[int],
D_diag: wp.array[int],
D_colind: wp.array[int],
# In:
qfrc: wp.array2d[float],
# Data out:
qacc_out: wp.array2d[float],
qLU_out: wp.array3d[float],
):
worldid = wp.tid()
NV = wp.static(nv)
# Phase 1: LU factorization (in-place on qLU_out)
for i in range(NV):
qacc_out[worldid, i] = float(D_rownnz[i])
# process diagonal elements from n-1 down to 0
for r_rev in range(NV):
i = NV - 1 - r_rev
rem_i = int(qacc_out[worldid, i])
rowadr_i = D_rowadr[i]
ii = rowadr_i + rem_i - 1
qacc_out[worldid, i] = float(rem_i - 1)
# cache diagonal element for row i
LUii = qLU_out[worldid, 0, ii]
# rows j above i (j < i), processed from i-1 down to 0
for c in range(i):
j = i - 1 - c
# get address of last remaining element of row j
rem_j = int(qacc_out[worldid, j])
rowadr_j = D_rowadr[j]
ji = rowadr_j + rem_j - 1
# process row j if (j,i) is non-zero
if D_colind[ji] == i:
# adjust remaining counter
rem_j = rem_j - 1
qacc_out[worldid, j] = float(rem_j)
# (j,i) = (j,i) / (i,i)
LUji = qLU_out[worldid, 0, ji] / LUii
qLU_out[worldid, 0, ji] = LUji
# (j,k) = (j,k) - (i,k) * (j,i) for k < i
icnt = rowadr_i
jcnt = rowadr_j
jend = rowadr_j + rem_j
while jcnt < jend:
col_i = D_colind[icnt]
col_j = D_colind[jcnt]
if col_i == col_j:
qLU_out[worldid, 0, jcnt] = qLU_out[worldid, 0, jcnt] - qLU_out[worldid, 0, icnt] * LUji
icnt = icnt + 1
jcnt = jcnt + 1
elif col_i > col_j:
jcnt = jcnt + 1
else:
icnt = icnt + 1
# Phase 2: LU solve (backward + forward substitution)
# Backward substitution: solve (U+I)*qacc = qfrc
for k_rev in range(NV):
i = NV - 1 - k_rev
diag_i = D_diag[i]
rowadr_i = D_rowadr[i]
d1 = diag_i + 1
nnz_upper = D_rownnz[i] - d1
acc = qfrc[worldid, i]
for j in range(nnz_upper):
adr_j = rowadr_i + d1 + j
col = D_colind[adr_j]
acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
qacc_out[worldid, i] = acc
# Forward substitution: solve L*qacc = qacc
for i in range(NV):
diag_i = D_diag[i]
rowadr_i = D_rowadr[i]
acc = qacc_out[worldid, i]
for j in range(diag_i):
adr_j = rowadr_i + j
col = D_colind[adr_j]
acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
qacc_out[worldid, i] = acc / qLU_out[worldid, 0, rowadr_i + diag_i]
return kernel
@event_scope
def factor_solve_lu(m: Model, d: Data, qLU: wp.array3d[float], qacc: wp.array2d[float], qfrc: wp.array2d[float]):
r"""Factorize and solve non-symmetric implicit system: qacc = A \\ qfrc.
qLU is overwritten in-place with the LU factors, then used to solve for qacc.
Args:
m: The model containing D-structure sparsity information.
d: The data object.
qLU: array containing the system matrix, overwritten with LU factors.
qacc: output array for the solution.
qfrc: input right-hand side.
"""
wp.launch(
_factor_solve_lu_sparse_fused(m.nv),
dim=(d.nworld,),
inputs=[m.D_rownnz, m.D_rowadr, m.D_diag, m.D_colind, qfrc],
outputs=[qacc, qLU],
)
@wp.kernel
def _subtree_vel_forward(
# Model:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+2227 -110
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+365 -15
View File
@@ -69,11 +69,11 @@ def mul_m_sparse(check_skip: bool):
@wp.kernel(module="unique")
def _mul_m_sparse(
# Model:
qM_mulm_rowadr: wp.array[int],
qM_mulm_col: wp.array[int],
qM_mulm_madr: wp.array[int],
M_mulm_rowadr: wp.array[int],
M_mulm_col: wp.array[int],
M_mulm_madr: wp.array[int],
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
vec: wp.array2d[float],
skip: wp.array[bool],
@@ -89,12 +89,12 @@ def mul_m_sparse(check_skip: bool):
# Gather all contributions (diagonal + off-diagonal)
acc = float(0.0)
start = qM_mulm_rowadr[dofid]
end = qM_mulm_rowadr[dofid + 1]
start = M_mulm_rowadr[dofid]
end = M_mulm_rowadr[dofid + 1]
for k in range(start, end):
col = qM_mulm_col[k]
madr = qM_mulm_madr[k]
acc += qM_in[worldid, 0, madr] * vec[worldid, col]
col = M_mulm_col[k]
madr = M_mulm_madr[k]
acc += M_in[worldid, 0, madr] * vec[worldid, col]
res[worldid, dofid] = acc
@@ -108,7 +108,7 @@ def mul_m_dense(nv: int, check_skip: bool):
@wp.kernel(module="unique")
def _mul_m_dense(
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
vec: wp.array2d[float],
skip: wp.array[bool],
@@ -123,7 +123,7 @@ def mul_m_dense(nv: int, check_skip: bool):
acc = float(0.0)
for j in range(wp.static(nv)):
acc += qM_in[worldid, i, j] * vec[worldid, j]
acc += M_in[worldid, i, j] * vec[worldid, j]
res[worldid, i] = acc
return _mul_m_dense
@@ -143,8 +143,8 @@ def mul_m(
Args:
m: The model containing kinematic and dynamic information (device).
d: The data object containing the current state and output arrays (device).
res: Result: qM @ vec.
vec: Input vector to multiply by qM.
res: Result: M @ vec.
vec: Input vector to multiply by M.
skip: Per-world bitmask to skip computing output.
M: Input matrix: M @ vec.
"""
@@ -152,13 +152,13 @@ def mul_m(
skip = skip or wp.empty(0, dtype=bool)
if M is None:
M = d.qM
M = d.M
if m.is_sparse:
wp.launch(
mul_m_sparse(check_skip),
dim=(d.nworld, m.nv),
inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
inputs=[m.M_mulm_rowadr, m.M_mulm_col, m.M_mulm_madr, M, vec, skip],
outputs=[res],
)
@@ -171,6 +171,340 @@ def mul_m(
)
@wp.kernel
def _mul_m_island_sparse(
# Model:
M_mulm_rowadr: wp.array[int],
M_mulm_col: wp.array[int],
M_mulm_madr: wp.array[int],
# Data in:
nidof_in: wp.array[int],
M_in: wp.array3d[float],
map_dof2idof_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
# In:
idof_islandid_in: wp.array2d[int],
vec: wp.array2d[float],
island_done_in: wp.array2d[bool],
check_skip: int,
# Out:
res: wp.array2d[float],
):
"""Sparse island mul_m for ALL islands in parallel."""
worldid, idofid = wp.tid()
nidof = nidof_in[worldid]
if idofid >= nidof:
return
islandid = idof_islandid_in[worldid, idofid]
if islandid < 0:
return
if check_skip:
if island_done_in[worldid, islandid]:
return
dof = map_idof2dof_in[worldid, idofid]
acc = float(0.0)
start = M_mulm_rowadr[dof]
end = M_mulm_rowadr[dof + 1]
for k in range(start, end):
col = M_mulm_col[k]
madr = M_mulm_madr[k]
col_idof = map_dof2idof_in[worldid, col]
# skip unconstrained DOFs
if col_idof < nidof:
acc += M_in[worldid, 0, madr] * vec[worldid, col_idof]
res[worldid, idofid] = acc
@wp.kernel
def _mul_m_island_dense(
# Model:
nv: int,
# Data in:
nidof_in: wp.array[int],
M_in: wp.array3d[float],
map_dof2idof_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
# In:
idof_islandid_in: wp.array2d[int],
vec: wp.array2d[float],
island_done_in: wp.array2d[bool],
check_skip: int,
# Out:
res: wp.array2d[float],
):
"""Dense island mul_m for ALL islands in parallel."""
worldid, idofid = wp.tid()
nidof = nidof_in[worldid]
if idofid >= nidof:
return
islandid = idof_islandid_in[worldid, idofid]
if islandid < 0:
return
if check_skip:
if island_done_in[worldid, islandid]:
return
dof = map_idof2dof_in[worldid, idofid]
acc = float(0.0)
for j in range(nv):
col_idof = map_dof2idof_in[worldid, j]
# skip unconstrained DOFs
if col_idof < nidof:
acc += M_in[worldid, dof, j] * vec[worldid, col_idof]
res[worldid, idofid] = acc
@event_scope
def mul_m_island(
m: Model,
d: Data,
res: wp.array2d[float],
vec: wp.array2d[float],
nidof: wp.array[int],
map_idof2dof: wp.array2d[int],
map_dof2idof: wp.array2d[int],
idof_islandid: wp.array2d[int],
island_done: Optional[wp.array] = None,
M: Optional[wp.array] = None,
):
"""Multiply island-local vectors by inertia matrix for all islands in parallel.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
res: Result: qM @ vec (island-local DOF order).
vec: Input vector (island-local DOF order).
nidof: Number of island DOFs per world.
map_idof2dof: Island-local DOF → global DOF map.
map_dof2idof: Global DOF → island-local DOF map.
idof_islandid: Island ID per island-local DOF.
island_done: Per-island done flags (nworld, ntree).
M: Optional mass matrix override.
"""
check_skip = int(island_done is not None)
island_done = island_done or wp.empty((0, 0), dtype=bool)
if M is None:
M = d.M
if m.is_sparse:
wp.launch(
_mul_m_island_sparse,
dim=(d.nworld, m.nv),
inputs=[
m.M_mulm_rowadr,
m.M_mulm_col,
m.M_mulm_madr,
nidof,
M,
map_dof2idof,
map_idof2dof,
idof_islandid,
vec,
island_done,
check_skip,
],
outputs=[res],
)
else:
wp.launch(
_mul_m_island_dense,
dim=(d.nworld, m.nv),
inputs=[
m.nv,
nidof,
M,
map_dof2idof,
map_idof2dof,
idof_islandid,
vec,
island_done,
check_skip,
],
outputs=[res],
)
@cache_kernel
def _solve_LD_sparse_island(nv: int, nlevels: int):
"""Sparse backsubstitution with island-local index remapping.
Same algorithm as _solve_LD_sparse_fused, but reads/writes x/y in
island-local DOF order. The L/D factorization stays in global DOF order.
"""
@wp.func_native(snippet="WP_TILE_SYNC();")
def _syncthreads():
pass
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Data in:
nidof_in: wp.array[int],
map_dof2idof_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
# In:
L: wp.array3d[float],
D: wp.array2d[float],
all_updates: wp.array[wp.vec3i],
level_offsets: wp.array[int],
y: wp.array2d[float],
# Out:
x_out: wp.array2d[float],
):
worldid, tid = wp.tid()
NLEVELS = wp.static(nlevels)
BLOCK_DIM = wp.block_dim()
nid = nidof_in[worldid]
# Copy y to x_out (island-local, only iterate up to nid)
for idof in range(tid, nid, BLOCK_DIM):
x_out[worldid, idof] = y[worldid, idof]
_syncthreads()
# Forward substitution
for level in range(NLEVELS):
level_idx = NLEVELS - 1 - level
level_offset = level_offsets[level_idx]
level_size = level_offsets[level_idx + 1] - level_offset
for u in range(tid, level_size, BLOCK_DIM):
update = all_updates[level_offset + u]
i, k, Madr_ki = update[0], update[1], update[2]
idof_i = map_dof2idof_in[worldid, i]
if idof_i < nid:
idof_k = map_dof2idof_in[worldid, k]
wp.atomic_sub(x_out[worldid], idof_i, L[worldid, 0, Madr_ki] * x_out[worldid, idof_k])
_syncthreads()
# Diagonal multiply (only iterate up to nid)
for idof in range(tid, nid, BLOCK_DIM):
dofid = map_idof2dof_in[worldid, idof]
x_out[worldid, idof] *= D[worldid, dofid]
_syncthreads()
# Backward substitution
for level in range(NLEVELS):
level_idx = level
level_offset = level_offsets[level_idx]
level_size = level_offsets[level_idx + 1] - level_offset
for u in range(tid, level_size, BLOCK_DIM):
update = all_updates[level_offset + u]
i, k, Madr_ki = update[0], update[1], update[2]
idof_k = map_dof2idof_in[worldid, k]
if idof_k < nid:
idof_i = map_dof2idof_in[worldid, i]
wp.atomic_sub(x_out[worldid], idof_k, L[worldid, 0, Madr_ki] * x_out[worldid, idof_i])
_syncthreads()
return kernel
@cache_kernel
def _tile_cholesky_solve_island(tile):
"""Dense Cholesky backsubstitution with island-local index remapping.
L is loaded from global DOF offsets (factorization unchanged).
y/x are loaded/stored at island-local DOF offsets via map_dof2idof.
"""
@wp.kernel(module="unique", enable_backward=False)
def cholesky_solve(
# Data in:
nidof_in: wp.array[int],
map_dof2idof_in: wp.array2d[int],
# In:
L: wp.array3d[float],
y: wp.array2d[float],
adr: wp.array[int],
# Out:
x: wp.array2d[float],
):
worldid, nodeid = wp.tid()
TILE_SIZE = wp.static(tile.size)
dofid = adr[nodeid]
idofid = map_dof2idof_in[worldid, dofid]
# Skip unconstrained trees (uniform branch — all threads in block agree)
if idofid >= nidof_in[worldid]:
return
# L stays in global order
L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
# y and x use island-local offsets
y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(idofid,))
x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
wp.tile_store(x[worldid], x_slice, offset=(idofid,))
return cholesky_solve
@event_scope
def solve_m_island(
m: Model,
d: Data,
res: wp.array2d[float],
vec: wp.array2d[float],
nidof: wp.array[int],
map_idof2dof: wp.array2d[int],
):
"""Compute res = M^{-1} @ vec for island-local DOFs.
Args:
m: Model.
d: Data.
res: Output in island-local DOF order.
vec: Input in island-local DOF order.
nidof: Number of island DOFs per world.
map_idof2dof: Island-local DOF -> global DOF map.
"""
if m.is_sparse:
nlevels = len(m.qLD_updates)
if wp.get_device().is_cuda:
dim_block = m.block_dim.solve_LD_sparse_fused
else:
dim_block = 1
wp.launch(
_solve_LD_sparse_island(m.nv, nlevels),
dim=(d.nworld, dim_block),
inputs=[
d.nidof,
d.map_dof2idof,
map_idof2dof,
d.qLD,
d.qLDiagInv,
m.qLD_all_updates,
m.qLD_level_offsets,
vec,
],
outputs=[res],
block_dim=dim_block,
)
else:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_solve_island(tile),
dim=(d.nworld, tile.adr.size),
inputs=[d.nidof, d.map_dof2idof, d.qLD, vec, tile.adr],
outputs=[res],
block_dim=m.block_dim.cholesky_solve,
)
@wp.kernel
def _apply_ft(
# Model:
@@ -604,11 +938,13 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
nbody: int,
neq: int,
nmocap: int,
nhistory: int,
# Data in:
time_in: wp.array[float],
qpos_in: wp.array2d[float],
qvel_in: wp.array2d[float],
act_in: wp.array2d[float],
history_in: wp.array2d[float],
qacc_warmstart_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
qfrc_applied_in: wp.array2d[float],
@@ -647,6 +983,10 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
for j in range(na):
state_out[worldid, adr + j] = act_in[worldid, j]
adr += na
elif element == State.HISTORY:
for j in range(nhistory):
state_out[worldid, adr + j] = history_in[worldid, j]
adr += nhistory
elif element == State.WARMSTART:
for j in range(nv):
state_out[worldid, adr + j] = qacc_warmstart_in[worldid, j]
@@ -700,10 +1040,12 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
m.nbody,
m.neq,
m.nmocap,
m.nhistory,
d.time,
d.qpos,
d.qvel,
d.act,
d.history,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
@@ -743,6 +1085,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
nbody: int,
neq: int,
nmocap: int,
nhistory: int,
# In:
sig_in: int,
active_in: wp.array[bool],
@@ -752,6 +1095,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
qpos_out: wp.array2d[float],
qvel_out: wp.array2d[float],
act_out: wp.array2d[float],
history_out: wp.array2d[float],
qacc_warmstart_out: wp.array2d[float],
ctrl_out: wp.array2d[float],
qfrc_applied_out: wp.array2d[float],
@@ -785,6 +1129,10 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
for j in range(na):
act_out[worldid, j] = state_in[worldid, adr + j]
adr += na
elif element == State.HISTORY:
for j in range(nhistory):
history_out[worldid, j] = state_in[worldid, adr + j]
adr += nhistory
elif element == State.WARMSTART:
for j in range(nv):
qacc_warmstart_out[worldid, j] = state_in[worldid, adr + j]
@@ -844,6 +1192,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
m.nbody,
m.neq,
m.nmocap,
m.nhistory,
int(sig),
active or wp.ones(d.nworld, dtype=bool),
state,
@@ -853,6 +1202,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
d.qpos,
d.qvel,
d.act,
d.history,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+258 -43
View File
@@ -17,18 +17,19 @@ import enum
from typing import Callable
import mujoco
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
import numpy as np
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
MJ_MINVAL = mujoco.mjMINVAL
MJ_MAXVAL = mujoco.mjMAXVAL
MJ_MINIMP = mujoco.mjMINIMP # minimum constraint impedance
MJ_MAXIMP = mujoco.mjMAXIMP # maximum constraint impedance
MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
MJ_MINMU = mujoco.mjMINMU # minimum friction
# True if MuJoCo >= 3.9.0 (new margin/gap semantics: includemargin = margin)
_NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0")
NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
# maximum size (by number of edges) of an horizon in EPA algorithm
MJ_MAX_EPAHORIZON = 24
# maximum average number of trianglarfaces EPA can insert at each iteration
@@ -47,6 +48,25 @@ class BlockDim:
"""Block dimension 'block_dim' settings for wp.launch_tiled.
TODO(team): experimental and may be removed
Attributes:
segmented_sort: segmented sort block dimension (collision_driver)
euler_dense: Euler dense block dimension (forward)
actuator_velocity: actuator velocity block dimension (forward)
ray: ray block dimension (ray)
contact_sort: contact sort block dimension (sensor)
energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
cholesky_factorize: Cholesky factorize block dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
"""
# collision_driver
@@ -131,13 +151,8 @@ class ProjectionType(enum.IntEnum):
ORTHOGRAPHIC: orthographic projection
"""
# TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
if hasattr(mujoco, "mjtProjection"):
PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
else:
PERSPECTIVE = 0
ORTHOGRAPHIC = 1
PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
class Stage(enum.IntEnum):
@@ -188,6 +203,7 @@ class DisableBit(enum.IntFlag):
EULERDAMP: implicit damping for Euler integration
NATIVECCD: native convex collision detection (ignored in MJWarp)
ISLAND: constraint islands
MULTICCD: disable multiple contacts with CCD
"""
CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
@@ -217,7 +233,6 @@ class EnableBit(enum.IntFlag):
Attributes:
ENERGY: energy computation
INVDISCRETE: discrete-time inverse dynamics
MULTICCD: multiple contacts with CCD
"""
ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
@@ -338,12 +353,13 @@ class IntegratorType(enum.IntEnum):
EULER: semi-implicit Euler
RK4: 4th-order Runge Kutta
IMPLICITFAST: implicit in velocity, no rne derivative
IMPLICIT: implicit in velocity, with rne derivative
"""
EULER = mujoco.mjtIntegrator.mjINT_EULER
RK4 = mujoco.mjtIntegrator.mjINT_RK4
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
# unsupported: IMPLICIT
IMPLICIT = mujoco.mjtIntegrator.mjINT_IMPLICIT
class GeomType(enum.IntEnum):
@@ -617,6 +633,7 @@ class State(enum.IntEnum):
QPOS: position
QVEL: velocity
ACT: actuator activation
HISTORY: delay/interval history buffers
WARMSTART: acceleration used for warmstart
CTRL: control
QFRC_APPLIED: applied generalized force
@@ -625,7 +642,7 @@ class State(enum.IntEnum):
MOCAP_POS: positions of mocap bodies
MOCAP_QUAT: orientations of mocap bodies
NSTATE: number of state elements
PHYSICS: QPOS | QVEL | ACT
PHYSICS: TIME | QPOS | QVEL | ACT | HISTORY
FULLPHYSICS: TIME | PHYSICS | PLUGIN
USER: CTRL | QFRC_APPLIED | XFRC_APPLIED | EQ_ACTIVE | MOCAP_POS | MOCAP_QUAT | USERDATA
INTEGRATION: FULLPHYSICS | USER | WARMSTART
@@ -635,6 +652,7 @@ class State(enum.IntEnum):
QPOS = mujoco.mjtState.mjSTATE_QPOS
QVEL = mujoco.mjtState.mjSTATE_QVEL
ACT = mujoco.mjtState.mjSTATE_ACT
HISTORY = mujoco.mjtState.mjSTATE_HISTORY
WARMSTART = mujoco.mjtState.mjSTATE_WARMSTART
CTRL = mujoco.mjtState.mjSTATE_CTRL
QFRC_APPLIED = mujoco.mjtState.mjSTATE_QFRC_APPLIED
@@ -643,7 +661,7 @@ class State(enum.IntEnum):
MOCAP_POS = mujoco.mjtState.mjSTATE_MOCAP_POS
MOCAP_QUAT = mujoco.mjtState.mjSTATE_MOCAP_QUAT
NSTATE = mujoco.mjtState.mjNSTATE
PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS
PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS # includes HISTORY
FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS
USER = mujoco.mjtState.mjSTATE_USER
INTEGRATION = mujoco.mjtState.mjSTATE_INTEGRATION
@@ -861,6 +879,7 @@ class Model:
ntree: number of kinematic trees
nM: number of non-zeros in sparse inertia matrix
nC: number of non-zeros in sparse body-dof matrix
nD: number of non-zeros in sparse derivative matrix
ngeom: number of geoms
nsite: number of sites
ncam: number of cameras
@@ -898,6 +917,7 @@ class Model:
nJmom: number of non-zeros in actuator_moment
ngravcomp: number of bodies with nonzero gravcomp
nsensordata: number of elements in sensor data vector
nhistory: number of history buffer entries
opt: physics options
stat: model statistics
qpos0: qpos values at default pose (*, nq)
@@ -1129,6 +1149,9 @@ class Model:
actuator_trnid: transmission id: joint, tendon, site (nu, 2)
actuator_actadr: first activation address; -1: stateless (nu,)
actuator_actnum: number of activation variables (nu,)
actuator_history: history buffer sizes (nu, 2)
actuator_historyadr: history buffer address (nu,)
actuator_delay: delay in seconds (nu,)
actuator_ctrllimited: is control limited (nu,)
actuator_forcelimited: is force limited (nu,)
actuator_actlimited: is activation limited (nu,)
@@ -1153,12 +1176,22 @@ class Model:
sensor_dim: number of scalar outputs (nsensor,)
sensor_adr: address in sensor array (nsensor,)
sensor_cutoff: cutoff for real and positive; 0: ignore (nsensor,)
sensor_history: history buffer sizes (nsensor, 2)
sensor_historyadr: history buffer address (nsensor,)
sensor_delay: delay in seconds (nsensor,)
sensor_interval: sensor interval and phase (nsensor, 2)
plugin: globally registered plugin slot number (nplugin,)
plugin_attr: config attributes of geom plugin (nplugin, _NPLUGINATTR)
M_rownnz: number of non-zeros in each row of qM (nv,)
M_rowadr: index of each row in qM (nv,)
M_colind: column indices of non-zeros in qM (nM,)
M_rownnz: number of non-zeros in each row of M (nv,)
M_rowadr: index of each row in M (nv,)
M_colind: column indices of non-zeros in M (nC,)
mapM2M: index mapping from M (legacy) to M (CSR) (nC)
D_rownnz: non-zeros per row in D-structure (nv,)
D_rowadr: row start addresses in D-structure (nv,)
D_diag: diagonal element index within each row (nv,)
D_colind: column indices in D-structure (nD,)
mapM2D: index mapping from M to D (nD,)
mapD2M: index mapping from D to M (nC,)
flex_vertflexid: flex id for each flex vertex (nflexvert,)
warp only fields:
@@ -1187,8 +1220,8 @@ class Model:
jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
jnt_limited_ball_adr: limited/ball jntadr
body_isdofancestor: precomputed mask of which DOFs affect each body
dof_tri_row: dof lower triangle row (used in solver)
dof_tri_col: dof lower triangle col (used in solver)
dof_tri_row: dof upper triangle row (used in solver)
dof_tri_col: dof upper triangle col (used in solver)
nxn_geom_pair: collision pair geom ids [-2, ngeom-1]
nxn_geom_pair_filtered: valid collision pair geom ids
[-1, ngeom - 1]
@@ -1243,15 +1276,21 @@ class Model:
sensor_rangefinder_bodyid: bodyid for rangefinder (nrangefinder,)
taxel_vertadr: tactile sensor vertex address (nsensortaxel,)
taxel_sensorid: address for tactile sensors
qM_tiles: tiling configuration
M_tiles: tiling configuration
qLD_updates: tuple of index triples for sparse factorization
qLD_all_updates: tuple of all levels concatenated
qLD_level_offsets: tuple of start offsets for each level
qM_fullm_i: sparse mass matrix addressing
qM_fullm_j: sparse mass matrix addressing
qM_mulm_rowadr: sparse matmul row pointers
qM_mulm_col: sparse matmul column indices
qM_mulm_madr: sparse matmul matrix addresses
M_fullm_i: sparse mass matrix addressing
M_fullm_j: sparse mass matrix addressing
M_elemid: (row, col) -> CSR madr addresses; -1 if not a chain ancestor
M_fullm_upper_i: upper-triangle row indices for solver h seeding
M_fullm_upper_j: upper-triangle column indices for solver h seeding
M_fullm_upper_elemid: source elemid into M_fullm_i/M_fullm_j
qD_fullm_i: D-structure row indices for RNE derivatives
qD_fullm_j: D-structure column indices for RNE derivatives
M_mulm_rowadr: sparse matmul row pointers
M_mulm_col: sparse matmul column indices
M_mulm_madr: sparse matmul matrix addresses
"""
nq: int
@@ -1264,6 +1303,7 @@ class Model:
ntree: int
nM: int
nC: int
nD: int
ngeom: int
nsite: int
ncam: int
@@ -1301,6 +1341,7 @@ class Model:
nJmom: int
ngravcomp: int
nsensordata: int
nhistory: int
opt: Option
stat: Statistic
qpos0: array("*", "nq", float)
@@ -1532,6 +1573,9 @@ class Model:
actuator_trnid: array("nu", wp.vec2i)
actuator_actadr: array("nu", int)
actuator_actnum: array("nu", int)
actuator_history: array("nu", wp.vec2i)
actuator_historyadr: array("nu", int)
actuator_delay: array("nu", float)
actuator_ctrllimited: array("nu", bool)
actuator_forcelimited: array("nu", bool)
actuator_actlimited: array("nu", bool)
@@ -1556,12 +1600,22 @@ class Model:
sensor_dim: array("nsensor", int)
sensor_adr: array("nsensor", int)
sensor_cutoff: array("nsensor", float)
sensor_history: array("nsensor", wp.vec2i)
sensor_historyadr: array("nsensor", int)
sensor_delay: array("nsensor", float)
sensor_interval: array("nsensor", wp.vec2)
plugin: array("nplugin", int)
plugin_attr: array("nplugin", vec_pluginattr)
M_rownnz: array("nv", int)
M_rowadr: array("nv", int)
M_colind: array("nC", int)
mapM2M: array("nC", int)
D_rownnz: array("nv", int)
D_rowadr: array("nv", int)
D_diag: array("nv", int)
D_colind: array("nD", int)
mapM2D: array("nD", int)
mapD2M: array("nC", int)
flex_vertflexid: array("nflexvert", int)
# warp only fields:
callback: Callback
@@ -1635,23 +1689,32 @@ class Model:
sensor_rangefinder_bodyid: array("nrangefinder", int)
taxel_vertadr: array("nsensortaxel", int)
taxel_sensorid: wp.array[int]
qM_tiles: tuple[TileSet, ...]
M_tiles: tuple[TileSet, ...]
qLD_updates: tuple[wp.array[wp.vec3i], ...]
qLD_all_updates: wp.array[wp.vec3i]
qLD_level_offsets: wp.array[int]
qM_fullm_i: wp.array[int]
qM_fullm_j: wp.array[int]
# TODO(team): Remove M_fullm_i/j and M_elemid by iterating the M CSR layout
# directly in the solver/derivative kernels
M_fullm_i: wp.array[int]
M_fullm_j: wp.array[int]
M_elemid: wp.array2d[int] # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
M_fullm_upper_i: wp.array[int]
M_fullm_upper_j: wp.array[int]
M_fullm_upper_elemid: wp.array[int]
qD_fullm_i: wp.array[int] # D-structure (full square) row indices for RNE derivatives
qD_fullm_j: wp.array[int] # D-structure (full square) column indices for RNE derivatives
# Gather-based sparse mul_m indices (thread per DOF, no atomics)
qM_mulm_rowadr: wp.array[int] # start address for each row [nv+1]
qM_mulm_col: wp.array[int] # column index to gather from
qM_mulm_madr: wp.array[int] # matrix address to read
M_mulm_rowadr: wp.array[int] # start address for each row [nv+1]
M_mulm_col: wp.array[int] # column index to gather from
M_mulm_madr: wp.array[int] # matrix address to read
class ContactType(enum.IntFlag):
"""Type of contact.
CONSTRAINT: contact for constraint solver.
SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO).
Attributes:
CONSTRAINT: contact for constraint solver
SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO)
"""
CONSTRAINT = 1
@@ -1673,6 +1736,8 @@ class Contact:
solimp: constraint solver impedance (naconmax, 5)
dim: contact space dimensionality: 1, 3, 4 or 6 (naconmax,)
geom: geom ids; -1 for flex (naconmax, 2)
flex: flex ids; -1 for geom (naconmax, 2)
vert: vertex ids for flex/mesh contact (naconmax, 2)
efc_address: address in efc; -1: not included (naconmax, nmaxpyramid)
worldid: world id (naconmax,)
type: ContactType (naconmax,)
@@ -1711,9 +1776,9 @@ class Constraint:
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
(nworld, 1, njmax_nnz) sparse
J: constraint Jacobian (nworld, njmax_pad, nv_pad) dense
(nworld, 1, njmax * nv) sparse
(nworld, 1, njmax_nnz) sparse
pos: constraint position (equality, contact) (nworld, njmax)
margin: inclusion margin (contact) (nworld, njmax)
D: constraint mass (nworld, njmax_pad)
@@ -1722,6 +1787,20 @@ class Constraint:
frictionloss: frictionloss (friction) (nworld, njmax)
force: constraint force in constraint space (nworld, njmax)
state: constraint state (nworld, njmax_pad)
island: island ID per constraint (nworld, njmax)
itype: island constraint type (nworld, njmax)
iid: island constraint id (nworld, njmax)
iJ_rownnz: island J_rownnz (nworld, njmax)
iJ_rowadr: island J_rowadr (nworld, njmax)
iJ_colind: island J_colind (nworld, 0, 0) dense
(nworld, 1, njmax_nnz) sparse
iJ: island J (nworld, njmax, nv) dense
(nworld, 1, njmax_nnz) sparse
iD: island constraint mass (nworld, njmax_pad)
iaref: island aref (nworld, njmax)
ifrictionloss: island frictionloss (nworld, njmax)
iforce: island force (nworld, njmax)
istate: island state (nworld, njmax_pad)
warp only fields:
Ma: M*qacc (nworld, nv)
Jqvel: J*qvel (nworld, njmax)
@@ -1729,8 +1808,8 @@ class Constraint:
type: array("nworld", "njmax", int)
id: array("nworld", "njmax", int)
J_rownnz: wp.array2d[int]
J_rowadr: wp.array2d[int]
J_rownnz: array("nworld", "njmax", int)
J_rowadr: array("nworld", "njmax", int)
J_colind: wp.array3d[int]
J: wp.array3d[float]
pos: array("nworld", "njmax", float)
@@ -1741,9 +1820,22 @@ class Constraint:
frictionloss: array("nworld", "njmax", float)
force: array("nworld", "njmax", float)
state: array("nworld", "njmax_pad", int)
island: array("nworld", "njmax", int)
Ma: array("nworld", "nv", float)
Jqvel: array("nworld", "njmax", float)
itype: array("nworld", "njmax", int)
iid: array("nworld", "njmax", int)
iJ_rownnz: array("nworld", "njmax", int)
iJ_rowadr: array("nworld", "njmax", int)
iJ_colind: wp.array3d[int]
iJ: wp.array3d[float]
iD: array("nworld", "njmax_pad", float)
iaref: array("nworld", "njmax", float)
ifrictionloss: array("nworld", "njmax", float)
iforce: array("nworld", "njmax", float)
istate: array("nworld", "njmax_pad", int)
@dataclasses.dataclass
class Data:
@@ -1756,11 +1848,13 @@ class Data:
nl: number of limit constraints (nworld,)
nefc: number of constraints (nworld,)
nisland: number of constraint islands (nworld,)
nidof: total DOFs in islands (nworld,)
time: simulation time (nworld,)
energy: potential, kinetic energy (nworld, 2)
qpos: position (nworld, nq)
qvel: velocity (nworld, nv)
act: actuator activation (nworld, na)
history: history buffer for delays (nworld, nhistory)
qacc_warmstart: acceleration used for warmstart (nworld, nv)
ctrl: control (nworld, nu)
qfrc_applied: applied generalized force (nworld, nv)
@@ -1791,7 +1885,7 @@ class Data:
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)
flexedge_length: flex edge lengths (nworld, nflexedge)
ten_wrapadr: start address of tendon's path (nworld, ntendon)
ten_wrapnum: number of wrap points in path (nworld, ntendon)
ten_J: tendon Jacobian (nworld, nJten)
@@ -1804,10 +1898,10 @@ class Data:
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
M: total inertia (nworld, nv_pad, nv_pad) if dense
(nworld, 1, nC) if sparse
qLD: upper Cholesky factorization (nworld, nv, nv) if dense
L'*D*L factorization of M (nworld, 1, nC) if sparse
qLDiagInv: 1/diag(D) (nworld, nv)
flexedge_velocity: flex edge velocities (nworld, nflexedge)
ten_velocity: tendon velocities (nworld, ntendon)
@@ -1822,6 +1916,7 @@ class Data:
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)
qLU: sparse LU factorization of (M - dt*qDeriv) (nworld, 1, nD)
actuator_force: actuator force in actuation space (nworld, nu)
qfrc_actuator: actuator force (nworld, nv)
qfrc_smooth: net unconstrained force (nworld, nv)
@@ -1836,6 +1931,23 @@ class Data:
contact: contact data
efc: constraint data
tree_island: island ID per tree (-1 if unconstrained) (nworld, ntree)
dof_island: island ID per DOF (-1 if unconstrained) (nworld, nv)
island_dofadr: island start address in dof vector (nworld, ntree)
island_nv: DOFs per island (nworld, ntree)
island_nefc: constraints per island (nworld, ntree)
island_ne: equality constraints per island (nworld, ntree)
island_nf: friction constraints per island (nworld, ntree)
island_efcadr: island start address in efc vector (nworld, ntree)
map_dof2idof: global DOF -> island-local DOF (nworld, nv)
map_idof2dof: island-local DOF -> global DOF (nworld, nv)
map_efc2iefc: global EFC -> island-local EFC (nworld, njmax)
map_iefc2efc: island-local EFC -> global EFC (nworld, njmax)
dof_islandid: island ID per island-DOF (nworld, nv)
efc_islandid: island ID per island-EFC (nworld, njmax)
iqacc: island-local qacc (nworld, nv)
iqacc_smooth: island-local qacc_smooth (nworld, nv)
iqfrc_smooth: island-local qfrc_smooth (nworld, nv)
iqfrc_constraint: island-local qfrc_constraint (nworld, nv)
warp only fields:
nworld: number of worlds
@@ -1854,11 +1966,13 @@ class Data:
nl: array("nworld", int)
nefc: array("nworld", int)
nisland: array("nworld", int)
nidof: array("nworld", int)
time: array("nworld", float)
energy: array("nworld", wp.vec2)
qpos: array("nworld", "nq", float)
qvel: array("nworld", "nv", float)
act: array("nworld", "na", float)
history: array("nworld", "nhistory", float)
qacc_warmstart: array("nworld", "nv", float)
ctrl: array("nworld", "nu", float)
qfrc_applied: array("nworld", "nv", float)
@@ -1902,7 +2016,7 @@ class Data:
moment_colind: array("nworld", "nJmom", int)
actuator_moment: array("nworld", "nJmom", float)
crb: array("nworld", "nbody", vec10)
qM: wp.array3d[float]
M: wp.array3d[float]
qLD: wp.array3d[float]
qLDiagInv: array("nworld", "nv", float)
flexedge_velocity: array("nworld", "nflexedge", float)
@@ -1918,6 +2032,7 @@ class Data:
qfrc_passive: array("nworld", "nv", float)
subtree_linvel: array("nworld", "nbody", wp.vec3)
subtree_angmom: array("nworld", "nbody", wp.vec3)
qLU: array("nworld", 1, "nD", float)
actuator_force: array("nworld", "nu", float)
qfrc_actuator: array("nworld", "nv", float)
qfrc_smooth: array("nworld", "nv", float)
@@ -1930,6 +2045,23 @@ class Data:
contact: Contact
efc: Constraint
tree_island: array("nworld", "ntree", int)
dof_island: array("nworld", "nv", int)
island_dofadr: array("nworld", "ntree", int)
island_nv: array("nworld", "ntree", int)
island_nefc: array("nworld", "ntree", int)
island_ne: array("nworld", "ntree", int)
island_nf: array("nworld", "ntree", int)
island_efcadr: array("nworld", "ntree", int)
map_dof2idof: array("nworld", "nv", int)
map_idof2dof: array("nworld", "nv", int)
map_efc2iefc: array("nworld", "njmax", int)
map_iefc2efc: array("nworld", "njmax", int)
dof_islandid: array("nworld", "nv", int)
efc_islandid: array("nworld", "njmax", int)
iqacc: wp.array2d[float]
iqacc_smooth: wp.array2d[float]
iqfrc_smooth: wp.array2d[float]
iqfrc_constraint: wp.array2d[float]
# warp only fields:
nworld: int
@@ -1942,6 +2074,77 @@ class Data:
ncollision: array(1, int)
@dataclasses.dataclass
class InverseContext:
"""Workspace arrays for inverse dynamics."""
Jaref: wp.array2d[float]
search_dot: wp.array[float]
gauss: wp.array[float]
cost: wp.array[float]
prev_cost: wp.array[float]
done: wp.array[bool]
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
@dataclasses.dataclass
class IslandSolverContext:
"""Workspace arrays for island constraint solver."""
# Re-ordered workspace arrays (sized per-dof / per-constraint)
Jaref: wp.array2d[float]
jv: wp.array2d[float]
search: wp.array2d[float]
mv: wp.array2d[float]
grad: wp.array2d[float]
Mgrad: wp.array2d[float]
prev_grad: wp.array2d[float]
prev_Mgrad: wp.array2d[float]
h: wp.array3d[float]
# Per-island solver scalars (nworld, ntree)
cost: wp.array2d[float]
prev_cost: wp.array2d[float]
gauss: wp.array2d[float]
search_dot: wp.array2d[float]
grad_dot: wp.array2d[float]
done: wp.array2d[bool] # per-island convergence
solver_niter: wp.array2d[int] # iterations per island
beta: wp.array2d[float]
alpha: wp.array2d[float]
Ma: wp.array2d[float] # island-local Ma (nworld, nv)
@dataclasses.dataclass
class SolverContext:
"""Workspace arrays for constraint solver."""
Jaref: wp.array2d[float]
search_dot: wp.array[float]
gauss: wp.array[float]
cost: wp.array[float]
prev_cost: wp.array[float]
done: wp.array[bool]
grad: wp.array2d[float]
grad_dot: wp.array[float]
Mgrad: wp.array2d[float]
search: wp.array2d[float]
mv: wp.array2d[float]
jv: wp.array2d[float]
quad: wp.array2d[wp.vec3]
quad_gauss: wp.array[wp.vec3]
alpha: wp.array[float]
prev_grad: wp.array2d[float]
prev_Mgrad: wp.array2d[float]
beta: wp.array[float]
h: wp.array3d[float]
hfactor: wp.array3d[float]
# Incremental Hessian update (Newton only)
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
@dataclasses.dataclass
class RenderContext:
"""Context for rendering.
@@ -1952,6 +2155,8 @@ class RenderContext:
cam_id_map: camera id map
use_textures: whether to use textures
use_shadows: whether to use shadows
use_ambient_lighting: whether to use ambient lighting
background_color: background color
use_precomputed_rays: whether to use precomputed rays
bvh_ngeom: number of geometries in the BVH
enabled_geom_ids: enabled geometry ids
@@ -1971,7 +2176,10 @@ class RenderContext:
flex_bvh_id: per-flex BVH ids
flex_group_root: per-flex group roots (nworld x n_flex_bvh)
flex_render_smooth: whether to render flex meshes smoothly
flex_dim: flex dimension per flex (1D/2D/3D)
bvh_nflexgeom: number of flex geometries in the BVH
flex_dim_np: flex dimension per flex (1D/2D/3D)
flex_geom_flexid: map from flex geom ID to flex ID
flex_geom_edgeid: map from flex geom ID to flex edge ID
bvh: scene BVH
bvh_id: scene BVH id
lower: lower bounds
@@ -1993,6 +2201,11 @@ class RenderContext:
render_skybox: whether to shade missed rays with the MuJoCo skybox texture
skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
enable_backface_culling: drop primitive ray hits whose normal faces away
from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
mesh-ray rule. When False, the renderer reports inner-surface hits, which
is faster but causes a camera placed inside a geom to render that geom's
back wall.
"""
nrender: int
@@ -2000,6 +2213,7 @@ class RenderContext:
cam_id_map: array("ncam", int)
use_textures: bool
use_shadows: bool
use_ambient_lighting: bool
background_color: wp.uint32
use_precomputed_rays: bool
render_skybox: bool
@@ -2045,3 +2259,4 @@ class RenderContext:
render_seg: array("ncam", bool)
znear: float
total_rays: int
enable_backface_culling: bool
mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
Vendored
+3 -3
View File
@@ -23,7 +23,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
"""Parse a version string into comparable components.
Dot-separated components form the version. Hyphen-separated suffixes (e.g.,
"-foo3") are treated as local build identifiers and stripped before
"-newton") are treated as local build identifiers and stripped before
parsing. 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
@@ -32,13 +32,13 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
Args:
version_str: Version string like "3.5.0", "3.5.0.dev869102767", or
"3.9.0-foo3".
"3.9.0-newton".
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.
"""
# Strip local build identifier (e.g., "3.9.0-foo3" -> "3.9.0")
# Strip local build identifier (e.g., "3.9.0-newton" -> "3.9.0")
version_str = version_str.split("-", 1)[0]
parts = version_str.split(".")
return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+2 -2
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name="mujoco-warp"
version = "3.8.0.1"
version = "3.9.0.1"
# TODO(team): create a distribution list
authors = [
{name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -31,7 +31,7 @@ dependencies = [
"etils[epath]",
"mujoco>=3.8.0",
"numpy",
"warp-lang>=1.12",
"warp-lang>=1.13",
]
[[tool.uv.index]]
mjx/mujoco/mjx/warp/__init__.py
+6
View File
@@ -23,6 +23,7 @@ if not typing.TYPE_CHECKING:
warp: Any = None
mujoco_warp: Any = None
mjwp_types: Any = None
mjwp_io: Any = None
WARP_INSTALLED: bool = False
# pylint: disable=g-import-not-at-top
@@ -35,6 +36,7 @@ if not typing.TYPE_CHECKING:
try:
from mujoco.mjx.third_party import mujoco_warp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.third_party.mujoco_warp._src import io as mjwp_io
except (ImportError, RuntimeError) as e:
print('Failed to import mujoco_warp:', e)
pass
@@ -67,7 +69,11 @@ else:
def BlockDim(self, *args, **kwargs): # pylint: disable=invalid-name
pass
class _MjwpIoStub:
ENABLE_ISLANDS = True
WARP_INSTALLED: bool = True
warp: Any = _WpStub()
mujoco_warp: Any = _MjwpStub()
mjwp_types: Any = _MjwpTypesStub()
mjwp_io: Any = _MjwpIoStub()
mjx/mujoco/mjx/warp/forward.py
+2581 -2118
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/forward_test.py
+2 -2
View File
@@ -176,8 +176,8 @@ class ForwardTest(parameterized.TestCase):
qm = np.zeros((m.nv, m.nv))
mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
# mjwarp adds padding to qM
tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
# mjwarp adds padding to M
tu.assert_eq(qm, dx._impl.M[: m.nv, : m.nv], 'M')
# qLD is fused in a cholesky factorize and solve, and not written to.
tu.assert_contact_eq(d, dx, worldid=i)
mjx/mujoco/mjx/warp/smooth_test.py
+2 -2
View File
@@ -112,7 +112,7 @@ class SmoothTest(parameterized.TestCase):
d = mujoco.MjData(m)
mx = mjx.put_model(m, impl='warp')
mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
rng = jax.random.PRNGKey(0)
dx = mjx.make_data(m, impl='warp')
@@ -188,7 +188,7 @@ class SmoothTest(parameterized.TestCase):
batch_size = 7
d = mujoco.MjData(m)
mx = mjx.put_model(m, impl='warp')
mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
worldids = jp.arange(batch_size)
dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
mjx/mujoco/mjx/warp/types.py
+201 -23
View File
@@ -78,6 +78,29 @@ class BlockDim:
"""Block dimension 'block_dim' settings for wp.launch_tiled.
TODO(team): experimental and may be removed
Attributes:
segmented_sort: segmented sort block dimension (collision_driver)
euler_dense: Euler dense block dimension (forward)
actuator_velocity: actuator velocity block dimension (forward)
ray: ray block dimension (ray)
contact_sort: contact sort block dimension (sensor)
energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
cholesky_factorize: Cholesky factorize block dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
cholesky_factorize_solve: Cholesky factorize and solve block dimension
(smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block
dimension (solver)
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension
(solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension
(solver)
linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
"""
actuator_velocity: int
@@ -134,10 +157,26 @@ class OptionWarp(PyTreeNode):
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
D_colind: np.ndarray
D_diag: np.ndarray
D_rowadr: np.ndarray
D_rownnz: np.ndarray
M_colind: np.ndarray
M_elemid: np.ndarray
M_fullm_i: np.ndarray
M_fullm_j: np.ndarray
M_fullm_upper_elemid: np.ndarray
M_fullm_upper_i: np.ndarray
M_fullm_upper_j: np.ndarray
M_mulm_col: np.ndarray
M_mulm_madr: np.ndarray
M_mulm_rowadr: np.ndarray
M_rowadr: np.ndarray
M_rownnz: np.ndarray
M_tiles: Tuple[TileSet, ...]
actuator_delay: np.ndarray
actuator_history: np.ndarray
actuator_historyadr: np.ndarray
actuator_trntype_body_adr: np.ndarray
block_dim: BlockDim
body_branch_start: np.ndarray
@@ -204,6 +243,8 @@ class ModelWarp(PyTreeNode):
light_active: jax.Array
light_bodyid: np.ndarray
light_targetbodyid: np.ndarray
mapD2M: np.ndarray
mapM2D: np.ndarray
mapM2M: np.ndarray
mat_texrepeat: jax.Array
max_ten_J_rownnz: int
@@ -249,22 +290,22 @@ class ModelWarp(PyTreeNode):
oct_coeff: np.ndarray
plugin: np.ndarray
plugin_attr: np.ndarray
qD_fullm_i: np.ndarray
qD_fullm_j: np.ndarray
qLD_all_updates: np.ndarray
qLD_level_offsets: np.ndarray
qLD_updates: Tuple[np.ndarray, ...]
qM_fullm_i: np.ndarray
qM_fullm_j: np.ndarray
qM_mulm_col: np.ndarray
qM_mulm_madr: np.ndarray
qM_mulm_rowadr: np.ndarray
qM_tiles: Tuple[TileSet, ...]
rangefinder_sensor_adr: np.ndarray
sensor_acc_adr: np.ndarray
sensor_adr_to_contact_adr: np.ndarray
sensor_collision_start_adr: np.ndarray
sensor_contact_adr: np.ndarray
sensor_delay: np.ndarray
sensor_e_kinetic: bool
sensor_e_potential: bool
sensor_history: np.ndarray
sensor_historyadr: np.ndarray
sensor_interval: np.ndarray
sensor_limitfrc_adr: np.ndarray
sensor_limitpos_adr: np.ndarray
sensor_limitvel_adr: np.ndarray
@@ -299,7 +340,7 @@ class ModelWarp(PyTreeNode):
class DataWarp(PyTreeNode):
"""Derived fields from Data."""
M: jax.Array
actuator_moment: jax.Array
actuator_velocity: jax.Array
cacc: jax.Array
@@ -323,6 +364,8 @@ class DataWarp(PyTreeNode):
contact__vert: jax.Array
contact__worldid: jax.Array
crb: jax.Array
dof_island: jax.Array
dof_islandid: jax.Array
efc__D: jax.Array
efc__J: jax.Array
efc__J_colind: jax.Array
@@ -333,19 +376,46 @@ class DataWarp(PyTreeNode):
efc__aref: jax.Array
efc__force: jax.Array
efc__frictionloss: jax.Array
efc__iD: jax.Array
efc__iJ: jax.Array
efc__iJ_colind: jax.Array
efc__iJ_rowadr: jax.Array
efc__iJ_rownnz: jax.Array
efc__iaref: jax.Array
efc__id: jax.Array
efc__iforce: jax.Array
efc__ifrictionloss: jax.Array
efc__iid: jax.Array
efc__island: jax.Array
efc__istate: jax.Array
efc__itype: jax.Array
efc__margin: jax.Array
efc__pos: jax.Array
efc__state: jax.Array
efc__type: jax.Array
efc__vel: jax.Array
efc_islandid: jax.Array
energy: jax.Array
flexedge_J: jax.Array
flexedge_length: jax.Array
flexedge_velocity: jax.Array
flexvert_xpos: jax.Array
iqacc: jax.Array
iqacc_smooth: jax.Array
iqfrc_constraint: jax.Array
iqfrc_smooth: jax.Array
island_dofadr: jax.Array
island_efcadr: jax.Array
island_ne: jax.Array
island_nefc: jax.Array
island_nf: jax.Array
island_nv: jax.Array
light_xdir: jax.Array
light_xpos: jax.Array
map_dof2idof: jax.Array
map_efc2iefc: jax.Array
map_idof2dof: jax.Array
map_iefc2efc: jax.Array
moment_colind: jax.Array
moment_rowadr: jax.Array
moment_rownnz: jax.Array
@@ -356,6 +426,7 @@ class DataWarp(PyTreeNode):
ne: jax.Array
nefc: jax.Array
nf: jax.Array
nidof: jax.Array
nisland: jax.Array
njmax: int
njmax_nnz: int
@@ -364,7 +435,7 @@ class DataWarp(PyTreeNode):
nworld: int
qLD: jax.Array
qLDiagInv: jax.Array
qM: jax.Array
qLU: jax.Array
qfrc_damper: jax.Array
qfrc_spring: jax.Array
solver_niter: jax.Array
@@ -401,6 +472,7 @@ DATA_NON_VMAP = {
'nacon',
'naconmax',
'ncollision',
'nidof',
'njmax',
'njmax_nnz',
'njmax_pad',
@@ -434,6 +506,7 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
_NDIM = {
'Data': {
'M': 3,
'act': 2,
'act_dot': 2,
'actuator_force': 2,
@@ -467,6 +540,8 @@ _NDIM = {
'crb': 3,
'ctrl': 2,
'cvel': 3,
'dof_island': 2,
'dof_islandid': 2,
'efc__D': 2,
'efc__J': 3,
'efc__J_colind': 3,
@@ -477,12 +552,25 @@ _NDIM = {
'efc__aref': 2,
'efc__force': 2,
'efc__frictionloss': 2,
'efc__iD': 2,
'efc__iJ': 3,
'efc__iJ_colind': 3,
'efc__iJ_rowadr': 2,
'efc__iJ_rownnz': 2,
'efc__iaref': 2,
'efc__id': 2,
'efc__iforce': 2,
'efc__ifrictionloss': 2,
'efc__iid': 2,
'efc__island': 2,
'efc__istate': 2,
'efc__itype': 2,
'efc__margin': 2,
'efc__pos': 2,
'efc__state': 2,
'efc__type': 2,
'efc__vel': 2,
'efc_islandid': 2,
'energy': 2,
'eq_active': 2,
'flexedge_J': 2,
@@ -491,8 +579,23 @@ _NDIM = {
'flexvert_xpos': 3,
'geom_xmat': 4,
'geom_xpos': 3,
'history': 2,
'iqacc': 2,
'iqacc_smooth': 2,
'iqfrc_constraint': 2,
'iqfrc_smooth': 2,
'island_dofadr': 2,
'island_efcadr': 2,
'island_ne': 2,
'island_nefc': 2,
'island_nf': 2,
'island_nv': 2,
'light_xdir': 3,
'light_xpos': 3,
'map_dof2idof': 2,
'map_efc2iefc': 2,
'map_idof2dof': 2,
'map_iefc2efc': 2,
'mocap_pos': 3,
'mocap_quat': 3,
'moment_colind': 2,
@@ -505,6 +608,7 @@ _NDIM = {
'ne': 1,
'nefc': 1,
'nf': 1,
'nidof': 1,
'nisland': 1,
'njmax': 0,
'njmax_nnz': 0,
@@ -513,7 +617,7 @@ _NDIM = {
'nworld': 0,
'qLD': 3,
'qLDiagInv': 2,
'qM': 3,
'qLU': 3,
'qacc': 2,
'qacc_smooth': 2,
'qacc_warmstart': 2,
@@ -556,9 +660,23 @@ _NDIM = {
'xquat': 3,
},
'Model': {
'D_colind': 1,
'D_diag': 1,
'D_rowadr': 1,
'D_rownnz': 1,
'M_colind': 1,
'M_elemid': 2,
'M_fullm_i': 1,
'M_fullm_j': 1,
'M_fullm_upper_elemid': 1,
'M_fullm_upper_i': 1,
'M_fullm_upper_j': 1,
'M_mulm_col': 1,
'M_mulm_madr': 1,
'M_mulm_rowadr': 1,
'M_rowadr': 1,
'M_rownnz': 1,
'M_tiles': -1,
'actuator_acc0': 2,
'actuator_actadr': 1,
'actuator_actearly': 1,
@@ -570,6 +688,7 @@ _NDIM = {
'actuator_cranklength': 2,
'actuator_ctrllimited': 1,
'actuator_ctrlrange': 3,
'actuator_delay': 1,
'actuator_dynprm': 3,
'actuator_dyntype': 1,
'actuator_forcelimited': 1,
@@ -577,6 +696,8 @@ _NDIM = {
'actuator_gainprm': 3,
'actuator_gaintype': 1,
'actuator_gear': 3,
'actuator_history': 2,
'actuator_historyadr': 1,
'actuator_lengthrange': 3,
'actuator_trnid': 2,
'actuator_trntype': 1,
@@ -769,6 +890,8 @@ _NDIM = {
'light_poscom0': 3,
'light_targetbodyid': 1,
'light_type': 2,
'mapD2M': 1,
'mapM2D': 1,
'mapM2M': 1,
'mat_rgba': 3,
'mat_texid': 3,
@@ -797,6 +920,7 @@ _NDIM = {
'mesh_vertnum': 1,
'mocap_bodyid': 1,
'nC': 0,
'nD': 0,
'nJfe': 0,
'nJmom': 0,
'nJten': 0,
@@ -821,6 +945,7 @@ _NDIM = {
'ngravcomp': 0,
'nhfield': 0,
'nhfielddata': 0,
'nhistory': 0,
'njnt': 0,
'nlight': 0,
'nmat': 0,
@@ -897,15 +1022,11 @@ _NDIM = {
'pair_solreffriction': 3,
'plugin': 1,
'plugin_attr': 2,
'qD_fullm_i': 1,
'qD_fullm_j': 1,
'qLD_all_updates': 2,
'qLD_level_offsets': 1,
'qLD_updates': -1,
'qM_fullm_i': 1,
'qM_fullm_j': 1,
'qM_mulm_col': 1,
'qM_mulm_madr': 1,
'qM_mulm_rowadr': 1,
'qM_tiles': -1,
'qpos0': 2,
'qpos_spring': 2,
'rangefinder_sensor_adr': 1,
@@ -916,9 +1037,13 @@ _NDIM = {
'sensor_contact_adr': 1,
'sensor_cutoff': 1,
'sensor_datatype': 1,
'sensor_delay': 1,
'sensor_dim': 1,
'sensor_e_kinetic': 0,
'sensor_e_potential': 0,
'sensor_history': 2,
'sensor_historyadr': 1,
'sensor_interval': 2,
'sensor_intprm': 2,
'sensor_limitfrc_adr': 1,
'sensor_limitpos_adr': 1,
@@ -1016,6 +1141,7 @@ _NDIM = {
}
_BATCH_DIM = {
'Data': {
'M': True,
'act': True,
'act_dot': True,
'actuator_force': True,
@@ -1049,6 +1175,8 @@ _BATCH_DIM = {
'crb': True,
'ctrl': True,
'cvel': True,
'dof_island': True,
'dof_islandid': True,
'efc__D': True,
'efc__J': True,
'efc__J_colind': True,
@@ -1059,12 +1187,25 @@ _BATCH_DIM = {
'efc__aref': True,
'efc__force': True,
'efc__frictionloss': True,
'efc__iD': True,
'efc__iJ': True,
'efc__iJ_colind': True,
'efc__iJ_rowadr': True,
'efc__iJ_rownnz': True,
'efc__iaref': True,
'efc__id': True,
'efc__iforce': True,
'efc__ifrictionloss': True,
'efc__iid': True,
'efc__island': True,
'efc__istate': True,
'efc__itype': True,
'efc__margin': True,
'efc__pos': True,
'efc__state': True,
'efc__type': True,
'efc__vel': True,
'efc_islandid': True,
'energy': True,
'eq_active': True,
'flexedge_J': True,
@@ -1073,8 +1214,23 @@ _BATCH_DIM = {
'flexvert_xpos': True,
'geom_xmat': True,
'geom_xpos': True,
'history': True,
'iqacc': True,
'iqacc_smooth': True,
'iqfrc_constraint': True,
'iqfrc_smooth': True,
'island_dofadr': True,
'island_efcadr': True,
'island_ne': True,
'island_nefc': True,
'island_nf': True,
'island_nv': True,
'light_xdir': True,
'light_xpos': True,
'map_dof2idof': True,
'map_efc2iefc': True,
'map_idof2dof': True,
'map_iefc2efc': True,
'mocap_pos': True,
'mocap_quat': True,
'moment_colind': True,
@@ -1087,6 +1243,7 @@ _BATCH_DIM = {
'ne': True,
'nefc': True,
'nf': True,
'nidof': False,
'nisland': True,
'njmax': False,
'njmax_nnz': False,
@@ -1095,7 +1252,7 @@ _BATCH_DIM = {
'nworld': False,
'qLD': True,
'qLDiagInv': True,
'qM': True,
'qLU': True,
'qacc': True,
'qacc_smooth': True,
'qacc_warmstart': True,
@@ -1138,9 +1295,23 @@ _BATCH_DIM = {
'xquat': True,
},
'Model': {
'D_colind': False,
'D_diag': False,
'D_rowadr': False,
'D_rownnz': False,
'M_colind': False,
'M_elemid': False,
'M_fullm_i': False,
'M_fullm_j': False,
'M_fullm_upper_elemid': False,
'M_fullm_upper_i': False,
'M_fullm_upper_j': False,
'M_mulm_col': False,
'M_mulm_madr': False,
'M_mulm_rowadr': False,
'M_rowadr': False,
'M_rownnz': False,
'M_tiles': False,
'actuator_acc0': True,
'actuator_actadr': False,
'actuator_actearly': False,
@@ -1152,6 +1323,7 @@ _BATCH_DIM = {
'actuator_cranklength': True,
'actuator_ctrllimited': False,
'actuator_ctrlrange': True,
'actuator_delay': False,
'actuator_dynprm': True,
'actuator_dyntype': False,
'actuator_forcelimited': False,
@@ -1159,6 +1331,8 @@ _BATCH_DIM = {
'actuator_gainprm': True,
'actuator_gaintype': False,
'actuator_gear': True,
'actuator_history': False,
'actuator_historyadr': False,
'actuator_lengthrange': True,
'actuator_trnid': False,
'actuator_trntype': False,
@@ -1351,6 +1525,8 @@ _BATCH_DIM = {
'light_poscom0': True,
'light_targetbodyid': False,
'light_type': True,
'mapD2M': False,
'mapM2D': False,
'mapM2M': False,
'mat_rgba': True,
'mat_texid': True,
@@ -1379,6 +1555,7 @@ _BATCH_DIM = {
'mesh_vertnum': False,
'mocap_bodyid': False,
'nC': False,
'nD': False,
'nJfe': False,
'nJmom': False,
'nJten': False,
@@ -1403,6 +1580,7 @@ _BATCH_DIM = {
'ngravcomp': False,
'nhfield': False,
'nhfielddata': False,
'nhistory': False,
'njnt': False,
'nlight': False,
'nmat': False,
@@ -1479,15 +1657,11 @@ _BATCH_DIM = {
'pair_solreffriction': True,
'plugin': False,
'plugin_attr': False,
'qD_fullm_i': False,
'qD_fullm_j': False,
'qLD_all_updates': False,
'qLD_level_offsets': False,
'qLD_updates': False,
'qM_fullm_i': False,
'qM_fullm_j': False,
'qM_mulm_col': False,
'qM_mulm_madr': False,
'qM_mulm_rowadr': False,
'qM_tiles': False,
'qpos0': True,
'qpos_spring': True,
'rangefinder_sensor_adr': False,
@@ -1498,9 +1672,13 @@ _BATCH_DIM = {
'sensor_contact_adr': False,
'sensor_cutoff': False,
'sensor_datatype': False,
'sensor_delay': False,
'sensor_dim': False,
'sensor_e_kinetic': False,
'sensor_e_potential': False,
'sensor_history': False,
'sensor_historyadr': False,
'sensor_interval': False,
'sensor_intprm': False,
'sensor_limitfrc_adr': False,
'sensor_limitpos_adr': False,