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: 868759523
Change-Id: I83a9f014840e66d7775a3684aa2e59ec5a077191
This commit is contained in:
Taylor Howell
2026-02-11 11:13:06 -08:00
committed by Copybara-Service
parent 883ba2134d
commit 94e834bf4f
33 changed files with 6783 additions and 2897 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/types.py
+2 -2
View File
@@ -871,10 +871,10 @@ class Model(PyTreeNode):
cam_poscom0: jax.Array
cam_pos0: jax.Array
cam_mat0: jax.Array
cam_fovy: np.ndarray
cam_fovy: jax.Array
cam_resolution: np.ndarray
cam_sensorsize: np.ndarray
cam_intrinsic: np.ndarray
cam_intrinsic: jax.Array
light_mode: np.ndarray
light_type: jax.Array
light_castshadow: jax.Array
mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
Vendored
+8
View File
@@ -28,6 +28,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model as Model
from mujoco.mjx.third_party.mujoco_warp._src.types import Data as Data
# isort: on
from mujoco.mjx.third_party.mujoco_warp._src.bvh import refit_bvh as refit_bvh
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import collision as collision
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import nxn_broadphase as nxn_broadphase
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import sap_broadphase as sap_broadphase
@@ -46,14 +47,19 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as runge
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.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
from mujoco.mjx.third_party.mujoco_warp._src.io import make_data as make_data
from mujoco.mjx.third_party.mujoco_warp._src.io import put_data as put_data
from mujoco.mjx.third_party.mujoco_warp._src.io import put_model as put_model
from mujoco.mjx.third_party.mujoco_warp._src.io import reset_data as reset_data
from mujoco.mjx.third_party.mujoco_warp._src.io import set_const as set_const
from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_0 as set_const_0
from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_fixed as set_const_fixed
from mujoco.mjx.third_party.mujoco_warp._src.passive import passive as passive
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray as ray
from mujoco.mjx.third_party.mujoco_warp._src.ray import rays as rays
from mujoco.mjx.third_party.mujoco_warp._src.render import render as render
from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_pos as energy_pos
from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_vel as energy_vel
from mujoco.mjx.third_party.mujoco_warp._src.sensor import sensor_acc as sensor_acc
@@ -75,6 +81,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.smooth import transmission as trans
from mujoco.mjx.third_party.mujoco_warp._src.solver import solve as solve
from mujoco.mjx.third_party.mujoco_warp._src.support import contact_force as contact_force
from mujoco.mjx.third_party.mujoco_warp._src.support import get_state as get_state
from mujoco.mjx.third_party.mujoco_warp._src.support import jac as jac
from mujoco.mjx.third_party.mujoco_warp._src.support import mul_m as mul_m
from mujoco.mjx.third_party.mujoco_warp._src.support import set_state as set_state
from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as xfrc_accumulate
@@ -92,6 +99,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType as GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType as IntegratorType
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType as JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Option as Option
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext as RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.types import SolverType as SolverType
from mujoco.mjx.third_party.mujoco_warp._src.types import State as State
from mujoco.mjx.third_party.mujoco_warp._src.types import Statistic as Statistic
mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
Vendored
+13 -4
View File
@@ -16,7 +16,7 @@
"""Utilities for benchmarking MuJoCo Warp."""
import time
from typing import Callable, Optional, Tuple
from typing import Callable, Tuple
import numpy as np
import warp as wp
@@ -24,6 +24,7 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import warp_util
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
@@ -87,10 +88,11 @@ def benchmark(
m: Model,
d: Data,
nstep: int,
ctrls: Optional[np.ndarray] = None,
ctrls: np.ndarray | None = None,
event_trace: bool = False,
measure_alloc: bool = False,
measure_solver_niter: bool = False,
render_context: RenderContext | None = None,
) -> Tuple[float, float, dict, list, list, list, int]:
"""Benchmark a function of Model and Data.
@@ -103,6 +105,7 @@ def benchmark(
event_trace: If True, time routines decorated with @event_scope.
measure_alloc: If True, record number of contacts and constraints.
measure_solver_niter: If True, record the number of solver iterations.
render_context: The render context to use for rendering.
Returns:
- Time to JIT fn.
@@ -120,8 +123,14 @@ def benchmark(
with warp_util.EventTracer(enabled=event_trace) as tracer:
# capture the whole function as a CUDA graph
jit_beg = time.perf_counter()
with wp.ScopedCapture() as capture:
fn(m, d)
if render_context is not None:
with wp.ScopedCapture() as capture:
fn(m, d, render_context)
else:
with wp.ScopedCapture() as capture:
fn(m, d)
jit_end = time.perf_counter()
jit_duration = jit_end - jit_beg
mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
Vendored
+1013
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+228 -249
View File
@@ -29,6 +29,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -38,7 +40,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
# TODO(team): improve compile time to enable backward pass
wp.set_module_options({"enable_backward": False})
@@ -46,42 +47,6 @@ wp.set_module_options({"enable_backward": False})
vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float)
mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
_CONVEX_COLLISION_PAIRS = [
(GeomType.HFIELD, GeomType.SPHERE),
(GeomType.HFIELD, GeomType.CAPSULE),
(GeomType.HFIELD, GeomType.ELLIPSOID),
(GeomType.HFIELD, GeomType.CYLINDER),
(GeomType.HFIELD, GeomType.BOX),
(GeomType.HFIELD, GeomType.MESH),
(GeomType.SPHERE, GeomType.ELLIPSOID),
(GeomType.SPHERE, GeomType.MESH),
(GeomType.CAPSULE, GeomType.ELLIPSOID),
(GeomType.CAPSULE, GeomType.CYLINDER),
(GeomType.CAPSULE, GeomType.MESH),
(GeomType.ELLIPSOID, GeomType.ELLIPSOID),
(GeomType.ELLIPSOID, GeomType.CYLINDER),
(GeomType.ELLIPSOID, GeomType.BOX),
(GeomType.ELLIPSOID, GeomType.MESH),
(GeomType.CYLINDER, GeomType.CYLINDER),
(GeomType.CYLINDER, GeomType.BOX),
(GeomType.CYLINDER, GeomType.MESH),
(GeomType.BOX, GeomType.MESH),
(GeomType.MESH, GeomType.MESH),
]
def _check_convex_collision_pairs():
prev_idx = -1
for pair in _CONVEX_COLLISION_PAIRS:
idx = upper_trid_index(len(GeomType), pair[0].value, pair[1].value)
if pair[1] < pair[0] or idx <= prev_idx:
return False
prev_idx = idx
return True
assert _check_convex_collision_pairs(), "_CONVEX_COLLISION_PAIRS is in invalid order."
@wp.func
def _hfield_filter(
@@ -91,6 +56,11 @@ def _hfield_filter(
geom_size: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graph: wp.array(dtype=int),
hfield_size: wp.array(dtype=wp.vec4),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -126,14 +96,14 @@ def _hfield_filter(
# box-sphere test: horizontal plane
for i in range(2):
if (size1[i] < pos[i] - r2 - margin) or (-size1[i] > pos[i] + r2 + margin):
return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
# box-sphere test: vertical direction
if size1[2] < pos[2] - r2 - margin: # up
return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
if -size1[3] > pos[2] + r2 + margin: # down
return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
mat2 = geom_xmat_in[worldid, g2]
mat = mat1T @ mat2
@@ -148,6 +118,15 @@ def _hfield_filter(
geomtype2 = geom_type[g2]
# load mesh vertex data for support function queries
if geomtype2 == GeomType.MESH:
dataid = geom_dataid[g2]
geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
geom2.vert = mesh_vert
geom2.graph = mesh_graph
# use support functions for tight AABB bounds
xmax = support(geom2, geomtype2, wp.vec3(1.0, 0.0, 0.0)).point[0]
xmin = support(geom2, geomtype2, wp.vec3(-1.0, 0.0, 0.0)).point[0]
@@ -165,7 +144,7 @@ def _hfield_filter(
or (zmin - margin > size1[2])
or (zmax + margin < -size1[3])
):
return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
else:
return False, xmin, xmax, ymin, ymax, zmin, zmax
@@ -176,11 +155,12 @@ def ccd_hfield_kernel_builder(
geomtype2: int,
gjk_iterations: int,
epa_iterations: int,
geomgeomid: int,
):
"""Kernel builder for heightfield CCD collisions (no multiccd args)."""
# runs convex collision on a set of geom pairs to recover contact info
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def ccd_hfield_kernel(
# Model:
opt_ccd_tolerance: wp.array(dtype=float),
@@ -196,15 +176,10 @@ def ccd_hfield_kernel_builder(
geom_friction: wp.array2d(dtype=wp.vec3),
geom_margin: wp.array2d(dtype=float),
geom_gap: wp.array2d(dtype=float),
hfield_adr: wp.array(dtype=int),
hfield_nrow: wp.array(dtype=int),
hfield_ncol: wp.array(dtype=int),
hfield_size: wp.array(dtype=wp.vec4),
hfield_data: wp.array(dtype=float),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graph: wp.array(dtype=int),
mesh_polynum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
@@ -215,6 +190,11 @@ def ccd_hfield_kernel_builder(
mesh_polymapadr: wp.array(dtype=int),
mesh_polymapnum: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
hfield_size: wp.array(dtype=wp.vec4),
hfield_nrow: wp.array(dtype=int),
hfield_ncol: wp.array(dtype=int),
hfield_adr: wp.array(dtype=int),
hfield_data: wp.array(dtype=float),
pair_dim: wp.array(dtype=int),
pair_solref: wp.array2d(dtype=wp.vec2),
pair_solreffriction: wp.array2d(dtype=wp.vec2),
@@ -223,24 +203,23 @@ def ccd_hfield_kernel_builder(
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# Data in:
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
naccdmax_in: int,
ncollision_in: wp.array(dtype=int),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
collision_worldid_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int),
# In:
epa_vert1_in: wp.array2d(dtype=wp.vec3),
epa_vert2_in: wp.array2d(dtype=wp.vec3),
epa_vert_index1_in: wp.array2d(dtype=int),
epa_vert_index2_in: wp.array2d(dtype=int),
epa_vert_in: wp.array2d(dtype=wp.vec3),
epa_vert_index_in: wp.array2d(dtype=int),
epa_face_in: wp.array2d(dtype=int),
epa_pr_in: wp.array2d(dtype=wp.vec3),
epa_norm2_in: wp.array2d(dtype=float),
epa_horizon_in: wp.array2d(dtype=int),
nccd_in: wp.array(dtype=int),
# Data out:
nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
@@ -254,27 +233,48 @@ def ccd_hfield_kernel_builder(
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
tid = wp.tid()
if tid >= ncollision_in[0]:
collisionid = wp.tid()
if collisionid >= ncollision_in[0]:
return
geoms = collision_pair_in[tid]
geoms = collision_pair_in[collisionid]
g1 = geoms[0]
g2 = geoms[1]
if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
return
worldid = collision_worldid_in[tid]
worldid = collision_worldid_in[collisionid]
# height field filter
no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter(
geom_type, geom_dataid, geom_size, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2
geom_type,
geom_dataid,
geom_size,
geom_rbound,
geom_margin,
mesh_vertadr,
mesh_vertnum,
mesh_graphadr,
mesh_vert,
mesh_graph,
hfield_size,
geom_xpos_in,
geom_xmat_in,
worldid,
g1,
g2,
)
if no_hf_collision:
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
_, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
geom_condim,
geom_priority,
@@ -293,7 +293,7 @@ def ccd_hfield_kernel_builder(
pair_friction,
collision_pair_in,
collision_pairid_in,
tid,
collisionid,
worldid,
)
@@ -365,26 +365,23 @@ def ccd_hfield_kernel_builder(
hfield_contact_dist = vec_maxconpair()
hfield_contact_pos = mat_maxconpair()
hfield_contact_normal = mat_maxconpair()
min_dist = float(wp.inf)
min_normal = wp.vec3(wp.inf, wp.inf, wp.inf)
min_pos = wp.vec3(wp.inf, wp.inf, wp.inf)
min_dist = float(MJ_MAXVAL)
min_normal = wp.vec3(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL)
min_pos = wp.vec3(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL)
min_id = int(-1)
# TODO(team): height field margin?
geom1.margin = margin
# geom1 margin added to z-axis of hfield prism top points below
geom2.margin = margin
# EPA memory
epa_vert1 = epa_vert1_in[tid]
epa_vert2 = epa_vert2_in[tid]
epa_vert_index1 = epa_vert_index1_in[tid]
epa_vert_index2 = epa_vert_index2_in[tid]
epa_face = epa_face_in[tid]
epa_pr = epa_pr_in[tid]
epa_norm2 = epa_norm2_in[tid]
epa_horizon = epa_horizon_in[tid]
epa_vert = epa_vert_in[ccdid]
epa_vert_index = epa_vert_index_in[ccdid]
epa_face = epa_face_in[ccdid]
epa_pr = epa_pr_in[ccdid]
epa_norm2 = epa_norm2_in[ccdid]
epa_horizon = epa_horizon_in[ccdid]
collision_pairid = collision_pairid_in[tid]
collision_pairid = collision_pairid_in[collisionid]
# process all prisms in subgrid
count = int(0)
@@ -439,11 +436,7 @@ def ccd_hfield_kernel_builder(
geom1.hfprism = prism
# prism center
x1 = geom1.pos
x1_ = wp.vec3(0.0, 0.0, 0.0)
for i in range(6):
x1_ += prism[i]
x1 += geom1.rot @ (x1_ / 6.0)
x1 = geom1.pos + geom1.rot @ (prism[0] + prism[1] + prism[2] + prism[3] + prism[4] + prism[5]) * wp.static(1.0 / 6.0)
dist, ncontact, w1, w2, idx = ccd(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
@@ -456,10 +449,8 @@ def ccd_hfield_kernel_builder(
geomtype2,
x1,
geom2.pos,
epa_vert1,
epa_vert2,
epa_vert_index1,
epa_vert_index2,
epa_vert,
epa_vert_index,
epa_face,
epa_pr,
epa_norm2,
@@ -529,13 +520,12 @@ def ccd_hfield_kernel_builder(
)
# TODO(team): routine for select subset of contacts
# TODO(team): if use_multiccd?
if wp.static(True):
MIN_DIST_TO_NEXT_CONTACT = 1.0e-3
# contact 1: furthest from minimum distance contact
id1 = int(-1)
dist1 = float(-wp.inf)
dist1 = float(-MJ_MAXVAL)
for i in range(count):
if i == min_id:
continue
@@ -589,7 +579,7 @@ def ccd_hfield_kernel_builder(
dist_min1 = wp.cross(min_normal, min_pos - pos1)
id2 = int(-1)
dist_12 = float(-wp.inf)
dist_12 = float(-MJ_MAXVAL)
for i in range(count):
if i == min_id or i == id1:
continue
@@ -644,7 +634,7 @@ def ccd_hfield_kernel_builder(
vec_12 = wp.cross(min_normal, pos1 - pos2)
id3 = int(-1)
dist3 = float(-wp.inf)
dist3 = float(-MJ_MAXVAL)
for i in range(count):
if i == min_id or i == id1 or i == id2:
continue
@@ -704,6 +694,7 @@ def ccd_kernel_builder(
gjk_iterations: int,
epa_iterations: int,
use_multiccd: bool,
geomgeomid: int,
):
"""Kernel builder for non-heightfield CCD collisions (no hfield args)."""
@@ -711,14 +702,11 @@ def ccd_kernel_builder(
def eval_ccd_write_contact(
# Model:
opt_ccd_tolerance: wp.array(dtype=float),
geom_type: wp.array(dtype=int),
# Data in:
naconmax_in: int,
# In:
epa_vert1_in: wp.array2d(dtype=wp.vec3),
epa_vert2_in: wp.array2d(dtype=wp.vec3),
epa_vert_index1_in: wp.array2d(dtype=int),
epa_vert_index2_in: wp.array2d(dtype=int),
epa_vert_in: wp.array2d(dtype=wp.vec3),
epa_vert_index_in: wp.array2d(dtype=int),
epa_face_in: wp.array2d(dtype=int),
epa_pr_in: wp.array2d(dtype=wp.vec3),
epa_norm2_in: wp.array2d(dtype=float),
@@ -738,7 +726,7 @@ def ccd_kernel_builder(
geom2: Geom,
geoms: wp.vec2i,
worldid: int,
tid: int,
ccdid: int,
margin: float,
gap: float,
condim: int,
@@ -748,7 +736,6 @@ def ccd_kernel_builder(
solimp: vec5,
x1: wp.vec3,
x2: wp.vec3,
count: int,
pairid: wp.vec2i,
# Data out:
contact_dist_out: wp.array(dtype=float),
@@ -771,12 +758,12 @@ def ccd_kernel_builder(
witness2 = mat43()
geom1.margin = margin
geom2.margin = margin
if pairid[1] >= 0:
# if collision sensor, set large cutoff to work with various sensor cutoff values
is_collision_sensor = pairid[1] >= 0
if is_collision_sensor:
cutoff = 1.0e32
else:
cutoff = 0.0
dist, ncontact, w1, w2, idx = ccd(
dist, ncollision, w1, w2, multiccd_idx = ccd(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
cutoff,
gjk_iterations,
@@ -787,14 +774,12 @@ def ccd_kernel_builder(
geomtype2,
x1,
x2,
epa_vert1_in[tid],
epa_vert2_in[tid],
epa_vert_index1_in[tid],
epa_vert_index2_in[tid],
epa_face_in[tid],
epa_pr_in[tid],
epa_norm2_in[tid],
epa_horizon_in[tid],
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 dist >= 0.0 and pairid[1] == -1:
@@ -803,30 +788,33 @@ def ccd_kernel_builder(
witness1[0] = w1
witness2[0] = w2
if wp.static(use_multiccd):
if (
geom1.margin == 0.0
and geom2.margin == 0.0
and (geomtype1 == GeomType.BOX or (geomtype1 == GeomType.MESH and geom1.mesh_polyadr > -1))
and (geomtype2 == GeomType.BOX or (geomtype2 == GeomType.MESH and geom2.mesh_polyadr > -1))
):
ncontact, witness1, witness2 = multicontact(
multiccd_polygon_in[tid],
multiccd_clipped_in[tid],
multiccd_pnormal_in[tid],
multiccd_pdist_in[tid],
multiccd_idx1_in[tid],
multiccd_idx2_in[tid],
multiccd_n1_in[tid],
multiccd_n2_in[tid],
multiccd_endvert_in[tid],
multiccd_face1_in[tid],
multiccd_face2_in[tid],
epa_vert1_in[tid],
epa_vert2_in[tid],
epa_vert_index1_in[tid],
epa_vert_index2_in[tid],
epa_face_in[tid, idx],
if wp.static(use_multiccd or (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)):
if wp.static(geomtype1 == GeomType.MESH):
# verify that geom1 mesh data is present for multicontact
if geom1.mesh_polyadr < 0:
multiccd_idx = -1
if wp.static(geomtype2 == GeomType.MESH):
# verify that geom2 mesh data is present for multicontact
if geom2.mesh_polyadr < 0:
multiccd_idx = -1
if multiccd_idx > -1:
ncollision, witness1, witness2 = multicontact(
multiccd_polygon_in[ccdid],
multiccd_clipped_in[ccdid],
multiccd_pnormal_in[ccdid],
multiccd_pdist_in[ccdid],
multiccd_idx1_in[ccdid],
multiccd_idx2_in[ccdid],
multiccd_n1_in[ccdid],
multiccd_n2_in[ccdid],
multiccd_endvert_in[ccdid],
multiccd_face1_in[ccdid],
multiccd_face2_in[ccdid],
epa_vert_in[ccdid],
epa_vert_index_in[ccdid],
epa_face_in[ccdid, multiccd_idx],
w1,
w2,
geom1,
@@ -835,7 +823,7 @@ def ccd_kernel_builder(
geomtype2,
)
for i in range(ncontact):
for i in range(ncollision):
points[i] = 0.5 * (witness1[i] + witness2[i])
normal = witness1[0] - witness2[0]
frame = make_frame(normal)
@@ -845,8 +833,9 @@ def ccd_kernel_builder(
frame *= -1.0
geoms = wp.vec2i(geoms[1], geoms[0])
for i in range(ncontact):
write_contact(
nactive = int(0) # number of contacts contributing to the physics
for i in range(ncollision):
active = write_contact(
naconmax_in,
i,
dist,
@@ -877,13 +866,12 @@ def ccd_kernel_builder(
contact_geomcollisionid_out,
nacon_out,
)
if count + (i + 1) >= MJ_MAXCONPAIR:
return i + 1
nactive += active
return ncontact
return nactive
# runs convex collision on a set of geom pairs to recover contact info (non-heightfield)
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def ccd_kernel(
# Model:
opt_ccd_tolerance: wp.array(dtype=float),
@@ -900,8 +888,8 @@ def ccd_kernel_builder(
geom_gap: wp.array2d(dtype=float),
mesh_vertadr: wp.array(dtype=int),
mesh_vertnum: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graphadr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_graph: wp.array(dtype=int),
mesh_polynum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
@@ -920,18 +908,17 @@ def ccd_kernel_builder(
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# Data in:
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
naccdmax_in: int,
ncollision_in: wp.array(dtype=int),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
collision_worldid_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int),
# In:
epa_vert1_in: wp.array2d(dtype=wp.vec3),
epa_vert2_in: wp.array2d(dtype=wp.vec3),
epa_vert_index1_in: wp.array2d(dtype=int),
epa_vert_index2_in: wp.array2d(dtype=int),
epa_vert_in: wp.array2d(dtype=wp.vec3),
epa_vert_index_in: wp.array2d(dtype=int),
epa_face_in: wp.array2d(dtype=int),
epa_pr_in: wp.array2d(dtype=wp.vec3),
epa_norm2_in: wp.array2d(dtype=float),
@@ -947,8 +934,8 @@ def ccd_kernel_builder(
multiccd_endvert_in: wp.array2d(dtype=wp.vec3),
multiccd_face1_in: wp.array2d(dtype=wp.vec3),
multiccd_face2_in: wp.array2d(dtype=wp.vec3),
nccd_in: wp.array(dtype=int),
# Data out:
nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
@@ -962,19 +949,25 @@ def ccd_kernel_builder(
contact_worldid_out: wp.array(dtype=int),
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
tid = wp.tid()
if tid >= ncollision_in[0]:
collisionid = wp.tid()
if collisionid >= ncollision_in[0]:
return
geoms = collision_pair_in[tid]
geoms = collision_pair_in[collisionid]
g1 = geoms[0]
g2 = geoms[1]
if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
return
worldid = collision_worldid_in[tid]
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(
geom_condim,
@@ -994,7 +987,7 @@ def ccd_kernel_builder(
pair_friction,
collision_pair_in,
collision_pairid_in,
tid,
collisionid,
worldid,
)
@@ -1024,12 +1017,9 @@ def ccd_kernel_builder(
eval_ccd_write_contact(
opt_ccd_tolerance,
geom_type,
naconmax_in,
epa_vert1_in,
epa_vert2_in,
epa_vert_index1_in,
epa_vert_index2_in,
epa_vert_in,
epa_vert_index_in,
epa_face_in,
epa_pr_in,
epa_norm2_in,
@@ -1049,7 +1039,7 @@ def ccd_kernel_builder(
geom2,
geoms,
worldid,
tid,
ccdid,
margin,
gap,
condim,
@@ -1059,8 +1049,7 @@ def ccd_kernel_builder(
solimp,
geom1.pos,
geom2.pos,
0,
collision_pairid_in[tid],
collision_pairid_in[collisionid],
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1080,37 +1069,8 @@ def ccd_kernel_builder(
return ccd_kernel
# Heightfield collision pairs handled by ccd_hfield_kernel_builder
_HFIELD_COLLISION_PAIRS = [
(GeomType.HFIELD, GeomType.SPHERE),
(GeomType.HFIELD, GeomType.CAPSULE),
(GeomType.HFIELD, GeomType.ELLIPSOID),
(GeomType.HFIELD, GeomType.CYLINDER),
(GeomType.HFIELD, GeomType.BOX),
(GeomType.HFIELD, GeomType.MESH),
]
# Non-heightfield collision pairs handled by ccd_kernel_builder
_NON_HFIELD_COLLISION_PAIRS = [
(GeomType.SPHERE, GeomType.ELLIPSOID),
(GeomType.SPHERE, GeomType.MESH),
(GeomType.CAPSULE, GeomType.ELLIPSOID),
(GeomType.CAPSULE, GeomType.CYLINDER),
(GeomType.CAPSULE, GeomType.MESH),
(GeomType.ELLIPSOID, GeomType.ELLIPSOID),
(GeomType.ELLIPSOID, GeomType.CYLINDER),
(GeomType.ELLIPSOID, GeomType.BOX),
(GeomType.ELLIPSOID, GeomType.MESH),
(GeomType.CYLINDER, GeomType.CYLINDER),
(GeomType.CYLINDER, GeomType.BOX),
(GeomType.CYLINDER, GeomType.MESH),
(GeomType.BOX, GeomType.MESH),
(GeomType.MESH, GeomType.MESH),
]
@event_scope
def convex_narrowphase(m: Model, d: Data):
def convex_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table: list[tuple[GeomType, GeomType]]):
"""Runs narrowphase collision detection for convex geom pairs.
This function handles collision detection for pairs of convex geometries that were
@@ -1126,40 +1086,56 @@ def convex_narrowphase(m: Model, d: Data):
computations for non-existent pair types.
"""
def _pair_count(p1: int, p2: int) -> int:
return m.geom_pair_type_count[upper_trid_index(len(GeomType), p1, p2)]
def _pair_count(p1: int, p2: int) -> Tuple[int, int]:
idx = upper_trid_index(len(GeomType), p1, p2)
return m.geom_pair_type_count[idx], idx
ncollision = sum(_pair_count(g[0].value, g[1].value)[0] for g in collision_table)
# no convex collisions, early return
if not any(_pair_count(g[0].value, g[1].value) for g in _CONVEX_COLLISION_PAIRS):
if ncollision == 0:
return
epa_iterations = m.opt.ccd_iterations
# compute nmaxpolygon and nmaxmeshdeg given the geom pairs for the model
nboxbox, _ = _pair_count(GeomType.BOX.value, GeomType.BOX.value)
if (GeomType.BOX, GeomType.BOX) not in collision_table:
nboxbox = 0
nboxmesh, _ = _pair_count(GeomType.BOX.value, GeomType.MESH.value)
nmeshmesh, _ = _pair_count(GeomType.MESH.value, GeomType.MESH.value)
epa_iterations = 16 if nboxbox == ncollision else m.opt.ccd_iterations
# set to true to enable multiccd
use_multiccd = m.opt.enableflags & EnableBit.MULTICCD
nmaxpolygon = m.nmaxpolygon if use_multiccd else 0
nmaxmeshdeg = m.nmaxmeshdeg if use_multiccd else 0
# epa_vert1: vertices in EPA polytope in geom 1 space
epa_vert1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
# epa_vert2: vertices in EPA polytope in geom 2 space
epa_vert2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
# epa_vert_index1: vertex indices in EPA polytope for geom 1
epa_vert_index1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int)
# epa_vert_index2: vertex indices in EPA polytope for geom 2 (naconmax, 5 + CCDiter)
epa_vert_index2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int)
# need at least 4 (square sides) if there's a box collision needing multiccd
nmaxpolygon = 4 if nboxbox > 0 else 0
nmaxmeshdeg = 3 if nboxbox > 0 else 0
# need to allocate more memory if there's meshes
if use_multiccd and nmeshmesh + nboxmesh > 0:
minval = 4 if nboxmesh else nmaxpolygon
nmaxpolygon = max(m.nmaxpolygon, minval)
nmaxmeshdeg = max(m.nmaxmeshdeg, 3)
# ccd collider count
nccd = wp.zeros(len(GeomType) * (len(GeomType) + 1) // 2, dtype=int)
# epa_vert: vertices in EPA polytope
epa_vert = wp.empty(shape=(d.naccdmax, 10 + 2 * epa_iterations), dtype=wp.vec3)
# epa_vert_index: vertex indices in EPA polytope
epa_vert_index = wp.empty(shape=(d.naccdmax, 10 + 2 * epa_iterations), dtype=int)
# epa_face: faces of polytope represented by three indices
epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
epa_face = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
# epa_pr: projection of origin on polytope faces
epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3)
epa_pr = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3)
# epa_norm2: epa_pr * epa_pr
epa_norm2 = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float)
epa_norm2 = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float)
# epa_horizon: index pair (i j) of edges on horizon
epa_horizon = wp.empty(shape=(d.naconmax, MJ_MAX_EPAHORIZON), dtype=int)
epa_horizon = wp.empty(shape=(d.naccdmax, MJ_MAX_EPAHORIZON), dtype=int)
# Contact outputs
contact_outputs = [
d.nacon,
d.contact.dist,
d.contact.pos,
d.contact.frame,
@@ -1173,15 +1149,17 @@ def convex_narrowphase(m: Model, d: Data):
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
]
# Launch heightfield collision kernels (no multiccd args, 72 args total)
for geom_pair in _HFIELD_COLLISION_PAIRS:
for geom_pair in collision_table:
g1 = geom_pair[0].value
g2 = geom_pair[1].value
if _pair_count(g1, g2):
count, geomgeomid = _pair_count(g1, g2)
if (g1 == GeomType.HFIELD or g2 == GeomType.HFIELD) and count:
wp.launch(
ccd_hfield_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations),
ccd_hfield_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, geomgeomid),
dim=d.naconmax,
inputs=[
m.opt.ccd_tolerance,
@@ -1197,15 +1175,10 @@ def convex_narrowphase(m: Model, d: Data):
m.geom_friction,
m.geom_margin,
m.geom_gap,
m.hfield_adr,
m.hfield_nrow,
m.hfield_ncol,
m.hfield_size,
m.hfield_data,
m.mesh_vertadr,
m.mesh_vertnum,
m.mesh_vert,
m.mesh_graphadr,
m.mesh_vert,
m.mesh_graph,
m.mesh_polynum,
m.mesh_polyadr,
@@ -1216,6 +1189,11 @@ def convex_narrowphase(m: Model, d: Data):
m.mesh_polymapadr,
m.mesh_polymapnum,
m.mesh_polymap,
m.hfield_size,
m.hfield_nrow,
m.hfield_ncol,
m.hfield_adr,
m.hfield_data,
m.pair_dim,
m.pair_solref,
m.pair_solreffriction,
@@ -1223,56 +1201,57 @@ def convex_narrowphase(m: Model, d: Data):
m.pair_margin,
m.pair_gap,
m.pair_friction,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.collision_pair,
d.collision_pairid,
d.collision_worldid,
d.naconmax,
d.naccdmax,
d.ncollision,
epa_vert1,
epa_vert2,
epa_vert_index1,
epa_vert_index2,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
epa_vert,
epa_vert_index,
epa_face,
epa_pr,
epa_norm2,
epa_horizon,
nccd,
],
outputs=contact_outputs,
)
# Allocate multiccd arrays only for non-heightfield collisions
# multiccd_polygon: clipped contact surface
multiccd_polygon = wp.empty(shape=(d.naconmax, 2 * nmaxpolygon), dtype=wp.vec3)
multiccd_polygon = wp.empty(shape=(d.naccdmax, 2 * nmaxpolygon), dtype=wp.vec3)
# multiccd_clipped: clipped contact surface (intermediate)
multiccd_clipped = wp.empty(shape=(d.naconmax, 2 * nmaxpolygon), dtype=wp.vec3)
multiccd_clipped = wp.empty(shape=(d.naccdmax, 2 * nmaxpolygon), dtype=wp.vec3)
# multiccd_pnormal: plane normal of clipping polygon
multiccd_pnormal = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=wp.vec3)
multiccd_pnormal = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=wp.vec3)
# multiccd_pdist: plane distance of clipping polygon
multiccd_pdist = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=float)
multiccd_pdist = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=float)
# multiccd_idx1: list of normal index candidates for Geom 1
multiccd_idx1 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=int)
multiccd_idx1 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=int)
# multiccd_idx2: list of normal index candidates for Geom 2
multiccd_idx2 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=int)
multiccd_idx2 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=int)
# multiccd_n1: list of normal candidates for Geom 1
multiccd_n1 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=wp.vec3)
multiccd_n1 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=wp.vec3)
# multiccd_n2: list of normal candidates for Geom 1
multiccd_n2 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=wp.vec3)
multiccd_n2 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=wp.vec3)
# multiccd_endvert: list of edge vertices candidates
multiccd_endvert = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=wp.vec3)
multiccd_endvert = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=wp.vec3)
# multiccd_face1: contact face
multiccd_face1 = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=wp.vec3)
multiccd_face1 = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=wp.vec3)
# multiccd_face2: contact face
multiccd_face2 = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=wp.vec3)
multiccd_face2 = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=wp.vec3)
# Launch non-heightfield collision kernels (no hfield args, 78 args total)
for geom_pair in _NON_HFIELD_COLLISION_PAIRS:
for geom_pair in collision_table:
g1 = geom_pair[0].value
g2 = geom_pair[1].value
if _pair_count(g1, g2):
count, geomgeomid = _pair_count(g1, g2)
if g1 != GeomType.HFIELD and g2 != GeomType.HFIELD and count:
wp.launch(
ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd),
ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd, geomgeomid),
dim=d.naconmax,
inputs=[
m.opt.ccd_tolerance,
@@ -1289,8 +1268,8 @@ def convex_narrowphase(m: Model, d: Data):
m.geom_gap,
m.mesh_vertadr,
m.mesh_vertnum,
m.mesh_vert,
m.mesh_graphadr,
m.mesh_vert,
m.mesh_graph,
m.mesh_polynum,
m.mesh_polyadr,
@@ -1308,17 +1287,16 @@ def convex_narrowphase(m: Model, d: Data):
m.pair_margin,
m.pair_gap,
m.pair_friction,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.collision_pair,
d.collision_pairid,
d.collision_worldid,
d.naconmax,
d.naccdmax,
d.ncollision,
epa_vert1,
epa_vert2,
epa_vert_index1,
epa_vert_index2,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
epa_vert,
epa_vert_index,
epa_face,
epa_pr,
epa_norm2,
@@ -1334,6 +1312,7 @@ def convex_narrowphase(m: Model, d: Data):
multiccd_endvert,
multiccd_face1,
multiccd_face2,
nccd,
],
outputs=contact_outputs,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+97 -36
View File
@@ -24,17 +24,65 @@ from mujoco.mjx.third_party.mujoco_warp._src.math import upper_tri_index
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
# Corresponding table to MuJoCo's mjCOLLISIONFUNC table in engine_collision_driver.c
MJ_COLLISION_TABLE = {
(GeomType.PLANE, GeomType.SPHERE): CollisionType.PRIMITIVE,
(GeomType.PLANE, GeomType.CAPSULE): CollisionType.PRIMITIVE,
(GeomType.PLANE, GeomType.ELLIPSOID): CollisionType.PRIMITIVE,
(GeomType.PLANE, GeomType.CYLINDER): CollisionType.PRIMITIVE,
(GeomType.PLANE, GeomType.BOX): CollisionType.PRIMITIVE,
(GeomType.PLANE, GeomType.MESH): CollisionType.PRIMITIVE,
(GeomType.HFIELD, GeomType.SPHERE): CollisionType.CONVEX,
(GeomType.HFIELD, GeomType.CAPSULE): CollisionType.CONVEX,
(GeomType.HFIELD, GeomType.ELLIPSOID): CollisionType.CONVEX,
(GeomType.HFIELD, GeomType.CYLINDER): CollisionType.CONVEX,
(GeomType.HFIELD, GeomType.BOX): CollisionType.CONVEX,
(GeomType.HFIELD, GeomType.MESH): CollisionType.CONVEX,
(GeomType.SPHERE, GeomType.SPHERE): CollisionType.PRIMITIVE,
(GeomType.SPHERE, GeomType.CAPSULE): CollisionType.PRIMITIVE,
(GeomType.SPHERE, GeomType.ELLIPSOID): CollisionType.CONVEX,
(GeomType.SPHERE, GeomType.CYLINDER): CollisionType.PRIMITIVE,
(GeomType.SPHERE, GeomType.BOX): CollisionType.PRIMITIVE,
(GeomType.SPHERE, GeomType.MESH): CollisionType.CONVEX,
(GeomType.CAPSULE, GeomType.CAPSULE): CollisionType.PRIMITIVE,
(GeomType.CAPSULE, GeomType.ELLIPSOID): CollisionType.CONVEX,
(GeomType.CAPSULE, GeomType.CYLINDER): CollisionType.CONVEX,
(GeomType.CAPSULE, GeomType.BOX): CollisionType.PRIMITIVE,
(GeomType.CAPSULE, GeomType.MESH): CollisionType.CONVEX,
(GeomType.ELLIPSOID, GeomType.ELLIPSOID): CollisionType.CONVEX,
(GeomType.ELLIPSOID, GeomType.CYLINDER): CollisionType.CONVEX,
(GeomType.ELLIPSOID, GeomType.BOX): CollisionType.CONVEX,
(GeomType.ELLIPSOID, GeomType.MESH): CollisionType.CONVEX,
(GeomType.CYLINDER, GeomType.CYLINDER): CollisionType.CONVEX,
(GeomType.CYLINDER, GeomType.BOX): CollisionType.CONVEX,
(GeomType.CYLINDER, GeomType.MESH): CollisionType.CONVEX,
(GeomType.BOX, GeomType.BOX): CollisionType.CONVEX, # overwritten by NATIVECCD disable flag
(GeomType.BOX, GeomType.MESH): CollisionType.CONVEX,
(GeomType.MESH, GeomType.MESH): CollisionType.CONVEX,
}
def create_collision_context(naconmax: int) -> CollisionContext:
"""Create a CollisionContext with allocated arrays."""
return CollisionContext(
collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
collision_pairid=wp.empty(naconmax, dtype=wp.vec2i),
collision_worldid=wp.empty(naconmax, dtype=int),
)
@wp.kernel
def _zero_nacon_ncollision(
@@ -293,10 +341,11 @@ def _add_geom_pair(
worldid: int,
nxnid: int,
# Data out:
ncollision_out: wp.array(dtype=int),
# Out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_pairid_out: wp.array(dtype=wp.vec2i),
collision_worldid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
pairid = wp.atomic_add(ncollision_out, 0, 1)
@@ -329,15 +378,15 @@ def _binary_search(values: wp.array(dtype=Any), value: Any, lower: int, upper: i
def _sap_project(opt_broadphase: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def sap_project(
# Model:
ngeom: int,
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
# Data in:
nworld_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
nworld_in: int,
# In:
direction_in: wp.vec3,
# Out:
@@ -402,7 +451,7 @@ def _sap_range(
@cache_kernel
def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
ngeom: int,
@@ -412,19 +461,20 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
geom_margin: wp.array2d(dtype=float),
nxn_pairid: wp.array(dtype=wp.vec2i),
# Data in:
nworld_in: int,
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
nworld_in: int,
naconmax_in: int,
# In:
sort_index_in: wp.array2d(dtype=int),
cumulative_sum_in: wp.array(dtype=int),
nsweep_in: int,
# Data out:
ncollision_out: wp.array(dtype=int),
# Out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_pairid_out: wp.array(dtype=wp.vec2i),
collision_worldid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
worldgeomid = wp.tid()
@@ -472,17 +522,17 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
geom2,
worldid,
idx,
ncollision_out,
collision_pair_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
return kernel
def _segmented_sort(tile_size: int):
@wp.kernel
@wp.kernel(module="unique")
def segmented_sort(
# In:
projection_lower_in: wp.array2d(dtype=float),
@@ -508,7 +558,7 @@ def _segmented_sort(tile_size: int):
@event_scope
def sap_broadphase(m: Model, d: Data):
def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
"""Runs broadphase collision detection using a sweep-and-prune (SAP) algorithm.
This method is more efficient than the N-squared approach for large numbers of
@@ -543,7 +593,7 @@ def sap_broadphase(m: Model, d: Data):
wp.launch(
kernel=_sap_project(m.opt.broadphase),
dim=(d.nworld, m.ngeom),
inputs=[m.ngeom, m.geom_rbound, m.geom_margin, d.nworld, d.geom_xpos, direction],
inputs=[m.ngeom, m.geom_rbound, m.geom_margin, d.geom_xpos, d.nworld, direction],
outputs=[
projection_lower.reshape((-1, m.ngeom)),
projection_upper,
@@ -588,21 +638,21 @@ def sap_broadphase(m: Model, d: Data):
m.geom_rbound,
m.geom_margin,
m.nxn_pairid,
d.nworld,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.nworld,
d.naconmax,
sort_index.reshape((-1, m.ngeom)),
cumulative_sum.reshape(-1),
nsweep,
],
outputs=[d.collision_pair, d.collision_pairid, d.collision_worldid, d.ncollision],
outputs=[d.ncollision, ctx.collision_pair, ctx.collision_pairid, ctx.collision_worldid],
)
@cache_kernel
def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
geom_type: wp.array(dtype=int),
@@ -612,14 +662,15 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
nxn_geom_pair: wp.array(dtype=wp.vec2i),
nxn_pairid: wp.array(dtype=wp.vec2i),
# Data in:
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
# Data out:
ncollision_out: wp.array(dtype=int),
# Out:
collision_pair_out: wp.array(dtype=wp.vec2i),
collision_pairid_out: wp.array(dtype=wp.vec2i),
collision_worldid_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
worldid, elementid = wp.tid()
@@ -641,17 +692,17 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
geom2,
worldid,
elementid,
ncollision_out,
collision_pair_out,
collision_pairid_out,
collision_worldid_out,
ncollision_out,
)
return kernel
@event_scope
def nxn_broadphase(m: Model, d: Data):
def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
"""Runs broadphase collision detection using a brute-force N-squared approach.
This function iterates through a pre-filtered list of all possible geometry pairs and
@@ -674,27 +725,34 @@ def nxn_broadphase(m: Model, d: Data):
m.geom_margin,
m.nxn_geom_pair_filtered,
m.nxn_pairid_filtered,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.naconmax,
],
outputs=[
d.collision_pair,
d.collision_pairid,
d.collision_worldid,
d.ncollision,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
],
)
def _narrowphase(m, d):
def _narrowphase(m: Model, d: Data, ctx: CollisionContext):
collision_table = MJ_COLLISION_TABLE
if m.opt.disableflags & DisableBit.NATIVECCD:
collision_table[(GeomType.BOX, GeomType.BOX)] = CollisionType.PRIMITIVE
convex_pairs = [key for key, value in collision_table.items() if value == CollisionType.CONVEX]
primitive_pairs = [key for key, value in collision_table.items() if value == CollisionType.PRIMITIVE]
# TODO(team): we should reject far-away contacts in the narrowphase instead of constraint
# partitioning because we can move some pressure of the atomics
convex_narrowphase(m, d)
primitive_narrowphase(m, d)
convex_narrowphase(m, d, ctx, convex_pairs)
primitive_narrowphase(m, d, ctx, primitive_pairs)
if m.has_sdf_geom:
sdf_narrowphase(m, d)
sdf_narrowphase(m, d, ctx)
@event_scope
@@ -715,15 +773,18 @@ def collision(m: Model, d: Data):
This function will do nothing except zero out arrays if collision detection is disabled
via `m.opt.disableflags` or if `d.nacon` is 0.
"""
# zero contact and collision counters
wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
if d.naconmax == 0 or m.opt.disableflags & (DisableBit.CONSTRAINT | DisableBit.CONTACT):
d.nacon.zero_()
return
if m.opt.broadphase == BroadphaseType.NXN:
nxn_broadphase(m, d)
else:
sap_broadphase(m, d)
ctx = create_collision_context(d.naconmax)
_narrowphase(m, d)
# zero counters
wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
if m.opt.broadphase == BroadphaseType.NXN:
nxn_broadphase(m, d, ctx)
else:
sap_broadphase(m, d, ctx)
_narrowphase(m, d, ctx)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+176 -177
View File
@@ -13,6 +13,7 @@
# limitations under the License.
# ==============================================================================
import math
from typing import Tuple
import warp as wp
@@ -30,9 +31,11 @@ FLOAT_MAX = 1e30
MINVAL = 1e-15
MIN_DIST = 1e-10
# TODO(kbayes): write out formulas to derive these constants
FACE_TOL = 0.99999872
EDGE_TOL = 0.00159999931
FACE_TOL = wp.static(math.cos(0.0016))
EDGE_TOL = wp.static(math.sin(0.0016))
# tolarance used by multicontact for intersecting a plane and a line segment
INTERSECT_TOL = 0.0000003
# Bit flags for face status in EPA polytope.
# Defined at module scope to avoid Warp's intermediate type issues with literals.
@@ -59,11 +62,9 @@ class GJKResult:
class Polytope:
status: int
# vertices in polytope
vert1: wp.array(dtype=wp.vec3)
vert2: wp.array(dtype=wp.vec3)
vert_index1: wp.array(dtype=int)
vert_index2: wp.array(dtype=int)
# vertices in polytope (packed geom1 followed by geom2)
vert: wp.array(dtype=wp.vec3)
vert_index: wp.array(dtype=int)
nvert: int
# faces in polytope
@@ -107,9 +108,9 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
tmp = wp.sign(local_dir)
res = wp.cw_mul(tmp, geom.size)
sp.point = geom.rot @ res + geom.pos
sp.vertex_index = wp.where(tmp[0] > 0, 1, 0)
sp.vertex_index += wp.where(tmp[1] > 0, 2, 0)
sp.vertex_index += wp.where(tmp[2] > 0, 4, 0)
sp.vertex_index = wp.where(tmp[0] > 0.0, 1, 0)
sp.vertex_index += wp.where(tmp[1] > 0.0, 2, 0)
sp.vertex_index += wp.where(tmp[2] > 0.0, 4, 0)
elif geomtype == GeomType.CAPSULE:
res = local_dir * geom.size[0]
# add cylinder contribution
@@ -177,7 +178,7 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
elif geomtype == GeomType.HFIELD:
max_dist = float(FLOAT_MIN)
# TODO(kbayes): Support edge prisms
sp.vertex_index = wp.where(dir[2] < 0, -2, -3)
sp.vertex_index = wp.where(dir[2] < 0.0, -2, -3)
for i in range(6):
vert = geom.hfprism[i]
dist = wp.dot(vert, dir)
@@ -197,7 +198,10 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float:
return 0.0
# compute witness point v
r, ret = _project_origin_plane(pt.vert1[v3] - pt.vert2[v3], pt.vert1[v2] - pt.vert2[v2], pt.vert1[v1] - pt.vert2[v1])
p1 = pt.vert[2 * v1] - pt.vert[2 * v1 + 1]
p2 = pt.vert[2 * v2] - pt.vert[2 * v2 + 1]
p3 = pt.vert[2 * v3] - pt.vert[2 * v3 + 1]
r, ret = _project_origin_plane(p3, p2, p1)
if ret:
return 0.0
@@ -214,13 +218,13 @@ def _epa_support(
pt: Polytope, idx: int, geom1: Geom, geom2: Geom, geom1_type: int, geom2_type: int, dir: wp.vec3
) -> Tuple[int, int]:
sp = support(geom1, geom1_type, dir)
pt.vert1[idx] = sp.point
pt.vert_index1[idx] = sp.vertex_index
pt.vert[2 * idx] = sp.point
pt.vert_index[2 * idx] = sp.vertex_index
index1 = sp.cached_index
sp = support(geom2, geom2_type, -dir)
pt.vert2[idx] = sp.point
pt.vert_index2[idx] = sp.vertex_index
pt.vert[2 * idx + 1] = sp.point
pt.vert_index[2 * idx + 1] = sp.vertex_index
index2 = sp.cached_index
return index1, index2
@@ -265,9 +269,9 @@ def _det3(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> float:
@wp.func
def _same_sign(a: float, b: float) -> int:
if a > 0 and b > 0:
if a > 0.0 and b > 0.0:
return 1
if a < 0 and b < 0:
if a < 0.0 and b < 0.0:
return -1
return 0
@@ -290,28 +294,25 @@ def _project_origin_plane(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> Tuple[wp.vec
n = wp.cross(diff32, diff21)
nv = wp.dot(n, v2)
nn = wp.dot(n, n)
if nn == 0:
if nn == 0.0:
return z, 1
if nv != 0 and nn > MINVAL:
v = (nv / nn) * n
return v, 0
if nv != 0.0 and nn > MINVAL:
return (nv / nn) * n, 0
# n = (v2 - v1) x (v3 - v1)
n = wp.cross(diff21, diff31)
nv = wp.dot(n, v1)
nn = wp.dot(n, n)
if nn == 0:
if nn == 0.0:
return z, 1
if nv != 0 and nn > MINVAL:
v = (nv / nn) * n
return v, 0
if nv != 0.0 and nn > MINVAL:
return (nv / nn) * n, 0
# n = (v1 - v3) x (v2 - v3)
n = wp.cross(diff31, diff32)
nv = wp.dot(n, v3)
nn = wp.dot(n, n)
v = (nv / nn) * n
return v, 0
return (nv / nn) * n, 0
@wp.func
@@ -610,14 +611,14 @@ def gjk(
simplex[n] = simplex1[n] - simplex2[n]
if cutoff == 0.0:
if wp.dot(x_k, simplex[n]) > 0:
if wp.dot(x_k, simplex[n]) > 0.0:
result = GJKResult()
result.dim = 0
result.dist = FLOAT_MAX
return result
elif cutoff < FLOAT_MAX:
vs = wp.dot(x_k, simplex[n])
if wp.dot(x_k, simplex[n]) > 0 and (vs * vs / xnorm) >= cutoff2:
if wp.dot(x_k, simplex[n]) > 0.0 and (vs * vs / xnorm) >= cutoff2:
result = GJKResult()
result.dim = 0
result.dist = FLOAT_MAX
@@ -635,7 +636,7 @@ def gjk(
# remove vertices from the simplex no longer needed
n = int(0)
for i in range(4):
if coordinates[i] == 0:
if coordinates[i] == 0.0:
continue
simplex[n] = simplex[i]
@@ -692,7 +693,7 @@ def _same_side(p0: wp.vec3, p1: wp.vec3, p2: wp.vec3, p3: wp.vec3) -> bool:
n = wp.cross(p1 - p0, p2 - p0)
dot1 = wp.dot(n, p3 - p0)
dot2 = wp.dot(n, -p0)
return (dot1 > 0 and dot2 > 0) or (dot1 < 0 and dot2 < 0)
return (dot1 > 0.0 and dot2 > 0.0) or (dot1 < 0.0 and dot2 < 0.0)
@wp.func
@@ -750,7 +751,7 @@ def _tri_point_intersect(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, p: wp.vec3) -> b
l2 = coordinates[1]
l3 = coordinates[2]
if l1 < 0 or l2 < 0 or l3 < 0:
if l1 < 0.0 or l2 < 0.0 or l3 < 0.0:
return False
pr = wp.vec3()
@@ -766,14 +767,14 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
# reset GJK simplex
simplex1 = mat43()
simplex1[0] = pt.vert1[v1]
simplex1[1] = pt.vert1[v2]
simplex1[2] = pt.vert1[v3]
simplex1[0] = pt.vert[2 * v1]
simplex1[1] = pt.vert[2 * v2]
simplex1[2] = pt.vert[2 * v3]
simplex2 = mat43()
simplex2[0] = pt.vert2[v1]
simplex2[1] = pt.vert2[v2]
simplex2[2] = pt.vert2[v3]
simplex2[0] = pt.vert[2 * v1 + 1]
simplex2[1] = pt.vert[2 * v2 + 1]
simplex2[2] = pt.vert[2 * v3 + 1]
simplex = mat43()
simplex[0] = simplex1[0] - simplex2[0]
@@ -781,14 +782,14 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
simplex[2] = simplex1[2] - simplex2[2]
simplex_index1 = wp.vec4i()
simplex_index1[0] = pt.vert_index1[v1]
simplex_index1[1] = pt.vert_index1[v2]
simplex_index1[2] = pt.vert_index1[v3]
simplex_index1[0] = pt.vert_index[2 * v1]
simplex_index1[1] = pt.vert_index[2 * v2]
simplex_index1[2] = pt.vert_index[2 * v3]
simplex_index2 = wp.vec4i()
simplex_index2[0] = pt.vert_index2[v1]
simplex_index2[1] = pt.vert_index2[v2]
simplex_index2[2] = pt.vert_index2[v3]
simplex_index2[0] = pt.vert_index[2 * v1 + 1]
simplex_index2[1] = pt.vert_index[2 * v2 + 1]
simplex_index2[2] = pt.vert_index[2 * v3 + 1]
result.simplex = simplex
result.simplex1 = simplex1
@@ -827,9 +828,9 @@ def _ray_triangle(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, v4: wp.vec3, v5: wp.vec
vol2 = _det3(v4 - v1, v5 - v1, v2 - v1)
vol3 = _det3(v5 - v1, v3 - v1, v2 - v1)
if vol1 >= 0 and vol2 >= 0 and vol3 >= 0:
if vol1 >= 0.0 and vol2 >= 0.0 and vol3 >= 0.0:
return 1
if vol1 <= 0 and vol2 <= 0 and vol3 <= 0:
if vol1 <= 0.0 and vol2 <= 0.0 and vol3 <= 0.0:
return -1
return 0
@@ -867,9 +868,9 @@ def _epa_witness(
) -> Tuple[wp.vec3, wp.vec3, float]:
face = _get_face_verts(pt.face[face_idx])
# compute affine coordinates for witness points on plane defined by face
v1 = pt.vert1[face[0]] - pt.vert2[face[0]]
v2 = pt.vert1[face[1]] - pt.vert2[face[1]]
v3 = pt.vert1[face[2]] - pt.vert2[face[2]]
v1 = pt.vert[2 * face[0]] - pt.vert[2 * face[0] + 1]
v2 = pt.vert[2 * face[1]] - pt.vert[2 * face[1] + 1]
v3 = pt.vert[2 * face[2]] - pt.vert[2 * face[2] + 1]
coordinates = _tri_affine_coord(v1, v2, v3, pt.face_pr[face_idx])
l1 = coordinates[0]
@@ -877,18 +878,18 @@ def _epa_witness(
l3 = coordinates[2]
# face on geom 2
v1 = pt.vert2[face[0]]
v2 = pt.vert2[face[1]]
v3 = pt.vert2[face[2]]
v1 = pt.vert[2 * face[0] + 1]
v2 = pt.vert[2 * face[1] + 1]
v3 = pt.vert[2 * face[2] + 1]
x2 = wp.vec3()
x2[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
x2[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
x2[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3
# correct witness points for hfield geoms
i1 = pt.vert_index1[face[0]]
i2 = pt.vert_index1[face[1]]
i3 = pt.vert_index1[face[2]]
i1 = pt.vert_index[2 * face[0]]
i2 = pt.vert_index[2 * face[1]]
i3 = pt.vert_index[2 * face[2]]
if geomtype1 == GeomType.HFIELD and (i1 != i2 or i1 != i3):
# TODO(kbayes): Fix case where geom2 is near bottom of height field or "extreme" prism heights
n = wp.vec3(0.0, 0.0, 1.0)
@@ -914,7 +915,7 @@ def _epa_witness(
x2 = sp.point
coordinates2 = _tri_affine_coord(a, b, c, x2)
if coordinates2[0] > 0 and coordinates2[1] > 0 and coordinates2[2] > 0:
if coordinates2[0] > 0.0 and coordinates2[1] > 0.0 and coordinates2[2] > 0.0:
x1 = coordinates2[0] * a + coordinates2[1] * b + coordinates2[2] * c
else:
p = c
@@ -924,9 +925,9 @@ def _epa_witness(
return x1, x2, -wp.norm_l2(x1 - x2)
# face on geom 1
v1 = pt.vert1[face[0]]
v2 = pt.vert1[face[1]]
v3 = pt.vert1[face[2]]
v1 = pt.vert[2 * face[0]]
v2 = pt.vert[2 * face[1]]
v3 = pt.vert[2 * face[2]]
x1 = wp.vec3()
x1[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
x1[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
@@ -971,17 +972,15 @@ def _polytope2(
d3 = R @ d2
# save vertices and get indices for each one
pt.vert1[0] = simplex1[0]
pt.vert1[1] = simplex1[1]
pt.vert[0] = simplex1[0]
pt.vert[1] = simplex2[0]
pt.vert[2] = simplex1[1]
pt.vert[3] = simplex2[1]
pt.vert_index1[0] = simplex_index1[0]
pt.vert_index1[1] = simplex_index1[1]
pt.vert2[0] = simplex2[0]
pt.vert2[1] = simplex2[1]
pt.vert_index2[0] = simplex_index2[0]
pt.vert_index2[1] = simplex_index2[1]
pt.vert_index[0] = simplex_index1[0]
pt.vert_index[1] = simplex_index2[0]
pt.vert_index[2] = simplex_index1[1]
pt.vert_index[3] = simplex_index2[1]
_epa_support(pt, 2, geom1, geom2, geomtype1, geomtype2, d1 / wp.norm_l2(d1))
_epa_support(pt, 3, geom1, geom2, geomtype1, geomtype2, d2 / wp.norm_l2(d2))
@@ -1013,9 +1012,9 @@ def _polytope2(
return pt, _replace_simplex3(pt, 1, 4, 3)
# check hexahedron is convex
v2 = pt.vert1[2] - pt.vert2[2]
v3 = pt.vert1[3] - pt.vert2[3]
v4 = pt.vert1[4] - pt.vert2[4]
v2 = pt.vert[4] - pt.vert[5]
v3 = pt.vert[6] - pt.vert[7]
v4 = pt.vert[8] - pt.vert[9]
if not _ray_triangle(simplex[0], simplex[1], v2, v3, v4):
pt.status = 1
return pt, GJKResult()
@@ -1049,21 +1048,19 @@ def _polytope3(
pt.status = 2
return pt
pt.vert1[0] = simplex1[0]
pt.vert1[1] = simplex1[1]
pt.vert1[2] = simplex1[2]
pt.vert[0] = simplex1[0]
pt.vert[1] = simplex2[0]
pt.vert[2] = simplex1[1]
pt.vert[3] = simplex2[1]
pt.vert[4] = simplex1[2]
pt.vert[5] = simplex2[2]
pt.vert_index1[0] = simplex_index1[0]
pt.vert_index1[1] = simplex_index1[1]
pt.vert_index1[2] = simplex_index1[2]
pt.vert2[0] = simplex2[0]
pt.vert2[1] = simplex2[1]
pt.vert2[2] = simplex2[2]
pt.vert_index2[0] = simplex_index2[0]
pt.vert_index2[1] = simplex_index2[1]
pt.vert_index2[2] = simplex_index2[2]
pt.vert_index[0] = simplex_index1[0]
pt.vert_index[1] = simplex_index2[0]
pt.vert_index[2] = simplex_index1[1]
pt.vert_index[3] = simplex_index2[1]
pt.vert_index[4] = simplex_index1[2]
pt.vert_index[5] = simplex_index2[2]
_epa_support(pt, 3, geom1, geom2, geomtype1, geomtype2, -n)
_epa_support(pt, 4, geom1, geom2, geomtype1, geomtype2, n)
@@ -1071,8 +1068,8 @@ def _polytope3(
v1 = simplex[0]
v2 = simplex[1]
v3 = simplex[2]
v4 = pt.vert1[3] - pt.vert2[3]
v5 = pt.vert1[4] - pt.vert2[4]
v4 = pt.vert[6] - pt.vert[7]
v5 = pt.vert[8] - pt.vert[9]
# check that v4 is not contained in the 2-simplex
if _tri_point_intersect(v1, v2, v3, v4):
@@ -1128,25 +1125,23 @@ def _polytope4(
simplex_index2: wp.vec4i,
) -> Tuple[Polytope, GJKResult]:
"""Create polytope for EPA given a 3-simplex from GJK."""
pt.vert1[0] = simplex1[0]
pt.vert1[1] = simplex1[1]
pt.vert1[2] = simplex1[2]
pt.vert1[3] = simplex1[3]
pt.vert[0] = simplex1[0]
pt.vert[1] = simplex2[0]
pt.vert[2] = simplex1[1]
pt.vert[3] = simplex2[1]
pt.vert[4] = simplex1[2]
pt.vert[5] = simplex2[2]
pt.vert[6] = simplex1[3]
pt.vert[7] = simplex2[3]
pt.vert_index1[0] = simplex_index1[0]
pt.vert_index1[1] = simplex_index1[1]
pt.vert_index1[2] = simplex_index1[2]
pt.vert_index1[3] = simplex_index1[3]
pt.vert2[0] = simplex2[0]
pt.vert2[1] = simplex2[1]
pt.vert2[2] = simplex2[2]
pt.vert2[3] = simplex2[3]
pt.vert_index2[0] = simplex_index2[0]
pt.vert_index2[1] = simplex_index2[1]
pt.vert_index2[2] = simplex_index2[2]
pt.vert_index2[3] = simplex_index2[3]
pt.vert_index[0] = simplex_index1[0]
pt.vert_index[1] = simplex_index2[0]
pt.vert_index[2] = simplex_index1[1]
pt.vert_index[3] = simplex_index2[1]
pt.vert_index[4] = simplex_index1[2]
pt.vert_index[5] = simplex_index2[2]
pt.vert_index[6] = simplex_index1[3]
pt.vert_index[7] = simplex_index2[3]
# if the origin is on a face, replace the 3-simplex with a 2-simplex
if _attach_face(pt, 0, 0, 1, 2) < MIN_DIST:
@@ -1249,7 +1244,7 @@ def _epa(
break
# check if lower bound is 0
if lower2 <= 0:
if lower2 <= 0.0:
break
# compute support point w from the closest face's normal
@@ -1257,7 +1252,7 @@ def _epa(
wi = pt.nvert
face_pr_normalized = pt.face_pr[idx] / lower
i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr_normalized)
w = pt.vert1[wi] - pt.vert2[wi]
w = pt.vert[2 * wi] - pt.vert[2 * wi + 1]
geom1.index = i1
geom2.index = i2
pt.nvert += 1
@@ -1275,7 +1270,7 @@ def _epa(
if is_discrete:
found_repeated = bool(False)
for i in range(pt.nvert - 1):
if pt.vert_index1[i] == pt.vert_index1[wi] and pt.vert_index2[i] == pt.vert_index2[wi]:
if pt.vert_index[2 * i] == pt.vert_index[2 * wi] and pt.vert_index[2 * i + 1] == pt.vert_index[2 * wi + 1]:
found_repeated = True
break
if found_repeated:
@@ -1313,7 +1308,7 @@ def _epa(
for i in range(pt.nhorizon):
edge = _get_edge(pt.horizon[i])
dist2 = _attach_face(pt, pt.nface, wi, edge[0], edge[1])
if dist2 == 0:
if dist2 == 0.0:
idx = -1
break
@@ -1348,72 +1343,66 @@ def _area4(a: wp.vec3, b: wp.vec3, c: wp.vec3, d: wp.vec3) -> float:
return 0.5 * wp.norm_l2(wp.cross(a - d, d - b) + wp.cross(b - c, c - a))
@wp.func
def _next(n: int, i: int) -> int:
"""Returns (i + 1) mod n for 0 <= i <= n - 1."""
return wp.where(i == n - 1, 0, i + 1)
@wp.func
def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i:
"""Returns the indices of a quadrilateral of maximum area in a convex polygon."""
b = _next(npolygon, 0)
c = _next(npolygon, b)
d = _next(npolygon, c)
"""Returns the indices of a quadrilateral of maximum area in a convex polygon (npolygon > 4)."""
b = int(1)
c = int(2)
d = int(3)
res = wp.vec4i(0, b, c, d)
m = _area4(polygon[0], polygon[b], polygon[c], polygon[d])
for a in range(npolygon):
while True:
m_next = _area4(polygon[a], polygon[b], polygon[c], polygon[_next(npolygon, d)])
m_next = _area4(polygon[a], polygon[b], polygon[c], polygon[(d + 1) % npolygon])
if m_next <= m:
break
m = m_next
d = _next(npolygon, d)
d = (d + 1) % npolygon
res = wp.vec4i(a, b, c, d)
while True:
m_next = _area4(polygon[a], polygon[b], polygon[_next(npolygon, c)], polygon[d])
m_next = _area4(polygon[a], polygon[b], polygon[(c + 1) % npolygon], polygon[d])
if m_next <= m:
break
m = m_next
c = _next(npolygon, c)
c = (c + 1) % npolygon
res = wp.vec4i(a, b, c, d)
while True:
m_next = _area4(polygon[a], polygon[_next(npolygon, b)], polygon[c], polygon[d])
m_next = _area4(polygon[a], polygon[(b + 1) % npolygon], polygon[c], polygon[d])
if m_next <= m:
break
m = m_next
b = _next(npolygon, b)
b = (b + 1) % npolygon
res = wp.vec4i(a, b, c, d)
if b == a:
b = _next(npolygon, b)
b = (b + 1) % npolygon
if c == b:
c = _next(npolygon, c)
c = (c + 1) % npolygon
if d == c:
d = _next(npolygon, d)
d = (d + 1) % npolygon
return res
# return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary
@wp.func
def _feature_dim(
face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3)
face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3), offset: int
) -> Tuple[int, wp.vec3i, wp.mat33]:
v1i = vert_index[face[0]]
v2i = vert_index[face[1]]
v3i = vert_index[face[2]]
v1i = vert_index[2 * face[0] + offset]
v2i = vert_index[2 * face[1] + offset]
v3i = vert_index[2 * face[2] + offset]
feature_index = wp.vec3i(v1i, v2i, v3i)
feature_vert = wp.mat33()
feature_vert[0] = vert[face[0]]
feature_vert[1] = vert[face[1]]
feature_vert[2] = vert[face[2]]
feature_vert[0] = vert[2 * face[0] + offset]
feature_vert[1] = vert[2 * face[1] + offset]
feature_vert[2] = vert[2 * face[2] + offset]
if v1i != v2i:
dim = wp.where(v3i == v1i or v3i == v2i, 2, 3)
return dim, feature_index, feature_vert
feature_index[1] = v3i
feature_vert[1] = vert[face[2]]
feature_vert[1] = vert[2 * face[2] + offset]
dim = wp.where(v1i != v3i, 2, 1)
return dim, feature_index, feature_vert
@@ -1675,7 +1664,7 @@ def _box_normals(
y = float((v1 & 2) and (v2 & 2)) - float(not (v1 & 2) and not (v2 & 2))
z = float((v1 & 4) and (v2 & 4)) - float(not (v1 & 4) and not (v2 & 4))
if x != 0.0:
normal_out[c] = mat @ wp.vec3(float(x), 0.0, 0.0)
normal_out[c] = mat @ wp.vec3(x, 0.0, 0.0)
index_out[c] = wp.where(x > 0.0, 0, 1)
c += 1
if y != 0.0:
@@ -1686,7 +1675,9 @@ def _box_normals(
normal_out[c] = mat @ wp.vec3(0.0, 0.0, z)
index_out[c] = wp.where(z > 0.0, 4, 5)
c += 1
if c == 2:
# c is 1 if edge is diagonal of a box face
# c is 2 if edge is an external edge of box
if c == 1 or c == 2:
return 2
return _box_normals2(mat, dir, normal_out, index_out)
@@ -1824,18 +1815,15 @@ def _halfspace(a: wp.vec3, n: wp.vec3, p: wp.vec3) -> bool:
@wp.func
def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> Tuple[float, wp.vec3]:
res = wp.vec3()
ab = b - a
temp = wp.dot(pn, ab)
if temp == 0.0:
return FLOAT_MAX, res # parallel; no intersection
t = (pd - wp.dot(pn, a)) / temp
if t >= 0.0 and t <= 1.0:
res[0] = a[0] + t * ab[0]
res[1] = a[1] + t * ab[1]
res[2] = a[2] + t * ab[2]
return t, res
def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> float:
"""Returns the parameter t where the line a + t(b - a) intersects the given plane."""
dot = wp.dot(pn, b - a)
# parallel; no intersection
if wp.abs(dot) < 1e-10:
return FLOAT_MAX
return (pd - wp.dot(pn, a)) / dot
# clip a polygon against another polygon
@@ -1901,9 +1889,10 @@ def _polygon_clip(
continue
# add new vertex to clipped polygon where PQ intersects the clipping edge
t, res = _plane_intersect(pn[e], pd[e], P, Q)
if t >= 0.0 and t <= 1.0:
clipped_out[nclipped] = res
t = _plane_intersect(pn[e], pd[e], P, Q)
if t > -INTERSECT_TOL and t < 1.0 + INTERSECT_TOL:
t = wp.clamp(t, 0.0, 1.0)
clipped_out[nclipped] = P + t * (Q - P)
nclipped += 1
# add Q as PQ is now back inside the clipping edge
@@ -1937,9 +1926,16 @@ def _polygon_clip(
@wp.func
def _set_edge(
vert1: wp.array(dtype=wp.vec3), vert2: wp.array(dtype=wp.vec3), start: int, end: int, face_out: wp.array(dtype=wp.vec3)
# In:
vert1: wp.array(dtype=wp.vec3),
vert2: wp.array(dtype=wp.vec3),
start: int,
end: int,
offset: int,
# Out:
face_out: wp.array(dtype=wp.vec3),
) -> int:
face_out[0] = vert1[start]
face_out[0] = vert1[2 * start + offset]
face_out[1] = vert2[end]
return 2
@@ -1959,10 +1955,8 @@ def multicontact(
endvert: wp.array(dtype=wp.vec3),
face1: wp.array(dtype=wp.vec3),
face2: wp.array(dtype=wp.vec3),
epa_vert1: wp.array(dtype=wp.vec3),
epa_vert2: wp.array(dtype=wp.vec3),
epa_vert_index1: wp.array(dtype=int),
epa_vert_index2: wp.array(dtype=int),
epa_vert: wp.array(dtype=wp.vec3),
epa_vert_index: wp.array(dtype=int),
epa_face: int,
x1: wp.vec3,
x2: wp.vec3,
@@ -1998,8 +1992,8 @@ def multicontact(
face = _get_face_verts(epa_face)
# get dimensions of features of geoms 1 and 2
nface1, feature_index1, feature_vertex1 = _feature_dim(face, epa_vert_index1, epa_vert1)
nface2, feature_index2, feature_vertex2 = _feature_dim(face, epa_vert_index2, epa_vert2)
nface1, feature_index1, feature_vertex1 = _feature_dim(face, epa_vert_index, epa_vert, 0)
nface2, feature_index2, feature_vertex2 = _feature_dim(face, epa_vert_index, epa_vert, 1)
dir = x2 - x1
dir_neg = -dir
@@ -2115,7 +2109,7 @@ def multicontact(
# recover geom1 matching edge or face
if is_edge_contact_geom1:
nface1 = _set_edge(epa_vert1, endvert, face[0], i, face1)
nface1 = _set_edge(epa_vert, endvert, face[0], i, 0, face1)
else:
ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i])
if geomtype1 == GeomType.BOX:
@@ -2136,7 +2130,7 @@ def multicontact(
# recover geom2 matching edge or face
if is_edge_contact_geom2:
nface2 = _set_edge(epa_vert2, endvert, face[0], i, face2)
nface2 = _set_edge(epa_vert, endvert, face[0], i, 1, face2)
else:
if geomtype2 == GeomType.BOX:
nface2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j], face2)
@@ -2197,12 +2191,12 @@ def _inflate(
c = geom1.hfprism[5]
coordinates = _tri_affine_coord(a, b, c, x2)
if coordinates[0] > 0 and coordinates[1] > 0 and coordinates[2] > 0:
if coordinates[0] > 0.0 and coordinates[1] > 0.0 and coordinates[2] > 0.0:
x1 = coordinates[0] * a + coordinates[1] * b + coordinates[2] * c
else:
p = c
p = wp.where(coordinates[1] > 0, b, p)
p = wp.where(coordinates[0] > 0, a, p)
p = wp.where(coordinates[1] > 0.0, b, p)
p = wp.where(coordinates[0] > 0.0, a, p)
x1 = x2 - wp.dot(x2 - p, n) * n
dist = -wp.norm_l2(x1 - x2)
return dist, x1, x2
@@ -2230,10 +2224,8 @@ def ccd(
geomtype2: int,
x_1: wp.vec3,
x_2: wp.vec3,
vert1: wp.array(dtype=wp.vec3),
vert2: wp.array(dtype=wp.vec3),
vert_index1: wp.array(dtype=int),
vert_index2: wp.array(dtype=int),
vert: wp.array(dtype=wp.vec3),
vert_index: wp.array(dtype=int),
face: wp.array(dtype=int),
face_pr: wp.array(dtype=wp.vec3),
face_norm2: wp.array(dtype=float),
@@ -2289,10 +2281,8 @@ def ccd(
pt.nface = 0
pt.nvert = 0
pt.nhorizon = 0
pt.vert1 = vert1
pt.vert2 = vert2
pt.vert_index1 = vert_index1
pt.vert_index2 = vert_index2
pt.vert = vert
pt.vert_index = vert_index
pt.face = face
pt.face_pr = face_pr
pt.face_norm2 = face_norm2
@@ -2358,4 +2348,13 @@ def ccd(
dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
if idx == -1:
return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
# multicontact not supported for margin
if geom1.margin != 0.0 or geom2.margin != 0.0:
idx = -1
# multicontact only supported for boxes and meshes
if (geomtype1 != GeomType.BOX and geomtype1 != GeomType.MESH) or (geomtype2 != GeomType.BOX and geomtype2 != GeomType.MESH):
idx = -1
return dist, 1, x1, x2, idx
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+19 -27
View File
@@ -32,8 +32,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -43,7 +45,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -173,13 +174,13 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
convex: Convex geometry object containing position, rotation, and mesh data.
Returns:
- Vector of contact distances (wp.inf for unpopulated contacts).
- Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
- Matrix of contact positions (one per row).
- Matrix of contact normal vectors (one per row).
"""
_HUGE_VAL = 1e6
contact_dist = wp.vec4(wp.inf)
contact_dist = wp.vec4(MJ_MAXVAL)
contact_pos = mat43()
contact_count = int(0)
@@ -426,12 +427,12 @@ def write_contact(
contact_type_out: wp.array(dtype=int),
contact_geomcollisionid_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
):
) -> int:
active = dist_in < margin_in
# skip contact and no collision sensor
if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1:
return
return 0
contact_type = 0
@@ -457,6 +458,8 @@ def write_contact(
contact_solimp_out[cid] = solimp_in
contact_type_out[cid] = contact_type
contact_geomcollisionid_out[cid] = id_
return int(active)
return 0
@wp.func
@@ -477,10 +480,9 @@ def contact_params(
pair_margin: wp.array2d(dtype=float),
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# Data in:
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
# In:
cid: int,
worldid: int,
):
@@ -1540,6 +1542,7 @@ def box_box_wrapper(
)
# Map of supported primitive collision functions
_PRIMITIVE_COLLISIONS = {
(GeomType.PLANE, GeomType.SPHERE): plane_sphere_wrapper,
(GeomType.PLANE, GeomType.CAPSULE): plane_capsule_wrapper,
@@ -1557,23 +1560,9 @@ _PRIMITIVE_COLLISIONS = {
}
# TODO(team): _check_collisions shared utility
def _check_primitive_collisions():
prev_idx = -1
for types in _PRIMITIVE_COLLISIONS.keys():
idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
if types[1] < types[0] or idx <= prev_idx:
return False
prev_idx = idx
return True
assert _check_primitive_collisions(), "_PRIMITIVE_COLLISIONS is in invalid order"
@cache_kernel
def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def primitive_narrowphase(
# Model:
geom_type: wp.array(dtype=int),
@@ -1612,10 +1601,11 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
ncollision_in: wp.array(dtype=int),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
collision_worldid_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int),
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
@@ -1730,7 +1720,7 @@ _PRIMITIVE_COLLISION_FUNC = []
@event_scope
def primitive_narrowphase(m: Model, d: Data):
def primitive_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table: list[tuple[GeomType, GeomType]]):
"""Runs collision detection on primitive geom pairs discovered during broadphase.
This function processes collision pairs involving primitive shapes that were
@@ -1749,6 +1739,8 @@ def primitive_narrowphase(m: Model, d: Data):
# for pair types without collisions, as well as updating the launch dimensions.
for types, func in _PRIMITIVE_COLLISIONS.items():
if types not in collision_table:
continue
idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
if m.geom_pair_type_count[idx] and types not in _PRIMITIVE_COLLISION_TYPES:
_PRIMITIVE_COLLISION_TYPES.append(types)
@@ -1793,10 +1785,10 @@ def primitive_narrowphase(m: Model, d: Data):
d.geom_xpos,
d.geom_xmat,
d.naconmax,
d.collision_pair,
d.collision_pairid,
d.collision_worldid,
d.ncollision,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
],
outputs=[
d.contact.dist,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py
Vendored
+9 -8
View File
@@ -18,6 +18,7 @@ from typing import Any, Tuple
import warp as wp
MJ_MINVAL = 1e-15
MJ_MAXVAL = 1e10
wp.set_module_options({"enable_backward": False})
@@ -411,13 +412,13 @@ def plane_box(
margin: Collision tolerance.
Returns:
- Vector of contact distances (wp.inf for unpopulated contacts).
- Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
- Matrix of contact positions (one per row).
- Contact normal vector.
"""
center_dist = wp.dot(box_pos - plane_pos, plane_normal)
dist = vec8f(wp.inf)
dist = vec8f(MJ_MAXVAL)
pos = mat83f()
# test all corners, pick bottom 4
@@ -540,7 +541,7 @@ def plane_cylinder(
- Matrix of contact normal vectors (one per row).
"""
# Initialize output matrices
contact_dist = wp.vec4(wp.inf)
contact_dist = wp.vec4(MJ_MAXVAL)
contact_pos = mat43f()
contact_count = 0
@@ -666,14 +667,14 @@ def box_box(
margin: Collision tolerance.
Returns:
- Vector of contact distances (wp.inf for unpopulated contacts).
- Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
- Matrix of contact positions (one per row).
- Matrix of contact normal vectors (one per row).
"""
# Initialize output matrices
contact_dist = vec8f()
for i in range(8):
contact_dist[i] = wp.inf
contact_dist[i] = MJ_MAXVAL
contact_pos = mat83f()
contact_normals = mat83f()
contact_count = 0
@@ -1176,7 +1177,7 @@ def capsule_box(
box_size: Half-extents of the box along each axis.
Returns:
- Vector of contact distances (wp.inf for unpopulated contacts).
- Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
- Matrix of contact positions (one per row).
- Matrix of contact normal vectors (one per row).
"""
@@ -1348,7 +1349,7 @@ def capsule_box(
c1 = wp.where((ee2 > 0) == w_neg, 1, 2)
if cltype == -4: # invalid type
return wp.vec2(wp.inf), mat23f(), mat23f()
return wp.vec2(MJ_MAXVAL), mat23f(), mat23f()
if cltype >= 0 and cltype // 3 != 1: # closest to a corner of the box
c1 = axisdir ^ clcorner
@@ -1479,7 +1480,7 @@ def capsule_box(
# collide with sphere using core function
dist2, pos2, normal2 = sphere_box(s2_pos_g, capsule_radius, box_pos, box_rot, box_size)
else:
dist2 = wp.inf
dist2 = MJ_MAXVAL
pos2 = wp.vec3()
normal2 = wp.vec3()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+13 -6
View File
@@ -22,6 +22,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_col
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -68,6 +69,7 @@ class MeshData:
data_id: int
pos: wp.vec3
mat: wp.mat33
size: wp.vec3
pnt: wp.vec3
vec: wp.vec3
valid: bool = False
@@ -375,6 +377,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
mesh_data.data_id,
mesh_data.pos,
mesh_data.mat,
mesh_data.size,
mesh_data.pnt,
mesh_data.vec,
)
@@ -388,6 +391,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
mesh_data.data_id,
mesh_data.pos,
mesh_data.mat,
mesh_data.size,
mesh_data.pnt,
-mesh_data.vec,
)
@@ -425,6 +429,7 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V
mesh_data.data_id,
mesh_data.pos,
mesh_data.mat,
mesh_data.size,
mesh_data.pnt,
mesh_data.vec,
)
@@ -664,11 +669,11 @@ def _sdf_narrowphase(
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
naconmax_in: int,
ncollision_in: wp.array(dtype=int),
# In:
collision_pair_in: wp.array(dtype=wp.vec2i),
collision_pairid_in: wp.array(dtype=wp.vec2i),
collision_worldid_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int),
# In:
sdf_initpoints: int,
sdf_iterations: int,
# Data out:
@@ -785,6 +790,7 @@ def _sdf_narrowphase(
mesh_data1.data_id = geom_dataid[g1]
mesh_data1.pos = geom1.pos
mesh_data1.mat = geom1.rot
mesh_data1.size = geom1.size
mesh_data1.pnt = wp.vec3(-1.0)
mesh_data1.vec = wp.vec3(0.0)
mesh_data1.valid = True
@@ -797,6 +803,7 @@ def _sdf_narrowphase(
mesh_data2.data_id = geom_dataid[g2]
mesh_data2.pos = geom2.pos
mesh_data2.mat = geom2.rot
mesh_data2.size = geom2.size
mesh_data2.pnt = wp.vec3(-1.0)
mesh_data2.vec = wp.vec3(0.0)
mesh_data2.valid = True
@@ -859,7 +866,7 @@ def _sdf_narrowphase(
@event_scope
def sdf_narrowphase(m: Model, d: Data):
def sdf_narrowphase(m: Model, d: Data, ctx: CollisionContext):
wp.launch(
_sdf_narrowphase,
dim=(m.opt.sdf_initpoints, d.naconmax),
@@ -909,10 +916,10 @@ def sdf_narrowphase(m: Model, d: Data):
d.geom_xpos,
d.geom_xmat,
d.naconmax,
d.collision_pair,
d.collision_pairid,
d.collision_worldid,
d.ncollision,
ctx.collision_pair,
ctx.collision_pairid,
ctx.collision_worldid,
m.opt.sdf_initpoints,
m.opt.sdf_iterations,
],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+174 -131
View File
@@ -34,21 +34,11 @@ def _zero_constraint_counts(
nf_out: wp.array(dtype=int),
nl_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
ne_connect_out: wp.array(dtype=int),
ne_weld_out: wp.array(dtype=int),
ne_jnt_out: wp.array(dtype=int),
ne_ten_out: wp.array(dtype=int),
ne_flex_out: wp.array(dtype=int),
):
worldid = wp.tid()
# Zero all constraint counters
ne_out[worldid] = 0
ne_connect_out[worldid] = 0
ne_weld_out[worldid] = 0
ne_jnt_out[worldid] = 0
ne_ten_out[worldid] = 0
ne_flex_out[worldid] = 0
nf_out[worldid] = 0
nl_out[worldid] = 0
nefc_out[worldid] = 0
@@ -155,6 +145,7 @@ def _efc_equality_connect(
# In:
refsafe_in: int,
# Data out:
ne_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
efc_type_out: wp.array2d(dtype=int),
efc_id_out: wp.array2d(dtype=int),
@@ -165,7 +156,6 @@ def _efc_equality_connect(
efc_vel_out: wp.array2d(dtype=float),
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
ne_connect_out: wp.array(dtype=int),
):
"""Calculates constraint rows for connect equality constraints."""
worldid, eqconnectid = wp.tid()
@@ -174,7 +164,7 @@ def _efc_equality_connect(
if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_connect_out, worldid, 3)
wp.atomic_add(ne_out, worldid, 3)
efcid = wp.atomic_add(nefc_out, worldid, 3)
if efcid + 3 >= njmax_in:
@@ -205,7 +195,7 @@ def _efc_equality_connect(
# compute Jacobian difference (opposite of contact: 0 - 1)
Jqvel = wp.vec3f(0.0, 0.0, 0.0)
for dofid in range(nv): # TODO: parallelize
jacp1, _ = support.jac(
jacp1, _ = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
@@ -216,7 +206,7 @@ def _efc_equality_connect(
dofid,
worldid,
)
jacp2, _ = support.jac(
jacp2, _ = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
@@ -293,6 +283,7 @@ def _efc_equality_joint(
# In:
refsafe_in: int,
# Data out:
ne_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
efc_type_out: wp.array2d(dtype=int),
efc_id_out: wp.array2d(dtype=int),
@@ -303,7 +294,6 @@ def _efc_equality_joint(
efc_vel_out: wp.array2d(dtype=float),
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
ne_jnt_out: wp.array(dtype=int),
):
worldid, eqjntid = wp.tid()
eqid = eq_jnt_adr[eqjntid]
@@ -311,7 +301,7 @@ def _efc_equality_joint(
if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_jnt_out, worldid, 1)
wp.atomic_add(ne_out, worldid, 1)
efcid = wp.atomic_add(nefc_out, worldid, 1)
if efcid >= njmax_in:
@@ -399,6 +389,7 @@ def _efc_equality_tendon(
# In:
refsafe_in: int,
# Data out:
ne_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
efc_type_out: wp.array2d(dtype=int),
efc_id_out: wp.array2d(dtype=int),
@@ -409,7 +400,6 @@ def _efc_equality_tendon(
efc_vel_out: wp.array2d(dtype=float),
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
ne_ten_out: wp.array(dtype=int),
):
worldid, eqtenid = wp.tid()
eqid = eq_ten_adr[eqtenid]
@@ -417,7 +407,7 @@ def _efc_equality_tendon(
if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_ten_out, worldid, 1)
wp.atomic_add(ne_out, worldid, 1)
efcid = wp.atomic_add(nefc_out, worldid, 1)
if efcid >= njmax_in:
@@ -494,6 +484,9 @@ def _efc_equality_flex(
opt_timestep: wp.array(dtype=float),
flexedge_length0: wp.array(dtype=float),
flexedge_invweight0: wp.array(dtype=float),
flexedge_J_rownnz: wp.array(dtype=int),
flexedge_J_rowadr: wp.array(dtype=int),
flexedge_J_colind: wp.array(dtype=int),
eq_solref: wp.array2d(dtype=wp.vec2),
eq_solimp: wp.array2d(dtype=vec5),
eq_flex_adr: wp.array(dtype=int),
@@ -505,6 +498,7 @@ def _efc_equality_flex(
# In:
refsafe_in: int,
# Data out:
ne_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
efc_type_out: wp.array2d(dtype=int),
efc_id_out: wp.array2d(dtype=int),
@@ -515,12 +509,11 @@ def _efc_equality_flex(
efc_vel_out: wp.array2d(dtype=float),
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
ne_flex_out: wp.array(dtype=int),
):
worldid, eqflexid, edgeid = wp.tid()
eqid = eq_flex_adr[eqflexid]
wp.atomic_add(ne_flex_out, worldid, 1)
wp.atomic_add(ne_out, worldid, 1)
efcid = wp.atomic_add(nefc_out, worldid, 1)
if efcid >= njmax_in:
@@ -531,10 +524,20 @@ def _efc_equality_flex(
solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid]
Jqvel = float(0.0)
# TODO(team): remove once efc.J is sparse
for i in range(nv):
J = flexedge_J_in[worldid, edgeid, i]
efc_J_out[worldid, efcid, i] = J
Jqvel += J * qvel_in[worldid, i]
efc_J_out[worldid, efcid, i] = 0.0
rownnz = flexedge_J_rownnz[edgeid]
rowadr = flexedge_J_rowadr[edgeid]
for i in range(rownnz):
sparseid = rowadr + i
colind = flexedge_J_colind[sparseid]
J = flexedge_J_in[worldid, 0, sparseid]
# TODO(team): sparse efc.J
efc_J_out[worldid, efcid, colind] = J
Jqvel += J * qvel_in[worldid, colind]
_update_efc_row(
worldid,
@@ -748,6 +751,7 @@ def _efc_equality_weld(
# In:
refsafe_in: int,
# Data out:
ne_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
efc_type_out: wp.array2d(dtype=int),
efc_id_out: wp.array2d(dtype=int),
@@ -758,7 +762,6 @@ def _efc_equality_weld(
efc_vel_out: wp.array2d(dtype=float),
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
ne_weld_out: wp.array(dtype=int),
):
worldid, eqweldid = wp.tid()
eqid = eq_wld_adr[eqweldid]
@@ -766,7 +769,7 @@ def _efc_equality_weld(
if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_weld_out, worldid, 6)
wp.atomic_add(ne_out, worldid, 6)
efcid = wp.atomic_add(nefc_out, worldid, 6)
if efcid + 6 >= njmax_in:
@@ -808,7 +811,7 @@ def _efc_equality_weld(
Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
for dofid in range(nv): # TODO: parallelize
jacp1, jacr1 = support.jac(
jacp1, jacr1 = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
@@ -819,7 +822,7 @@ def _efc_equality_weld(
dofid,
worldid,
)
jacp2, jacr2 = support.jac(
jacp2, jacr2 = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
@@ -1215,8 +1218,12 @@ def _efc_contact_pyramidal(
opt_impratio_invsqrt: wp.array(dtype=float),
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
body_weldid: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int),
body_dofadr: wp.array(dtype=int),
body_invweight0: wp.array2d(dtype=wp.vec2),
dof_bodyid: wp.array(dtype=int),
dof_parentid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
# Data in:
qvel_in: wp.array2d(dtype=float),
@@ -1302,49 +1309,73 @@ def _efc_contact_pyramidal(
invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt
Jqvel = float(0.0)
for i in range(nv):
J = float(0.0)
Ji = float(0.0)
jac1p, jac1r = support.jac(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
con_pos,
body1,
i,
worldid,
)
jac2p, jac2r = support.jac(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
con_pos,
body2,
i,
worldid,
)
jacp_dif = jac2p - jac1p
for xyz in range(3):
J += frame[0, xyz] * jacp_dif[xyz]
# skip fixed bodies
body1 = body_weldid[body1]
body2 = body_weldid[body2]
da1 = body_dofadr[body1] + body_dofnum[body1] - 1
da2 = body_dofadr[body2] + body_dofnum[body2] - 1
da = wp.max(da1, da2)
for dofid in range(nv - 1, -1, -1):
if dofid == da:
# TODO(team): contact_jacobian
jac1p, jac1r = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
con_pos,
body1,
dofid,
worldid,
)
jac2p, jac2r = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
con_pos,
body2,
dofid,
worldid,
)
J = float(0.0)
Ji = float(0.0)
if condim > 1:
dimid2 = dimid / 2 + 1
for xyz in range(3):
jacp_dif = jac2p[xyz] - jac1p[xyz]
J += frame[0, xyz] * jacp_dif
if condim > 1:
if dimid2 < 3:
Ji += frame[dimid2, xyz] * jacp_dif
else:
Ji += frame[dimid2 - 3, xyz] * (jac2r[xyz] - jac1r[xyz])
if condim > 1:
if dimid2 < 3:
Ji += frame[dimid2, xyz] * jacp_dif[xyz]
if dimid % 2 == 0:
J += Ji * frii
else:
Ji += frame[dimid2 - 3, xyz] * (jac2r[xyz] - jac1r[xyz])
J -= Ji * frii
if condim > 1:
if dimid % 2 == 0:
J += Ji * frii
else:
J -= Ji * frii
efc_J_out[worldid, efcid, dofid] = J
Jqvel += J * qvel_in[worldid, dofid]
efc_J_out[worldid, efcid, i] = J
Jqvel += J * qvel_in[worldid, i]
# Advance tree pointers and recompute da for next iteration
if da1 == da:
da1 = dof_parentid[da1]
if da2 == da:
da2 = dof_parentid[da2]
da = wp.max(da1, da2)
else:
efc_J_out[worldid, efcid, dofid] = 0.0
if condim == 1:
efc_type = ConstraintType.CONTACT_FRICTIONLESS
@@ -1385,8 +1416,12 @@ def _efc_contact_elliptic(
opt_impratio_invsqrt: wp.array(dtype=float),
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
body_weldid: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int),
body_dofadr: wp.array(dtype=int),
body_invweight0: wp.array2d(dtype=wp.vec2),
dof_bodyid: wp.array(dtype=int),
dof_parentid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
# Data in:
qvel_in: wp.array2d(dtype=float),
@@ -1450,49 +1485,69 @@ def _efc_contact_elliptic(
impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
contact_efc_address_out[conid, dimid] = efcid
con_pos = pos_in[conid]
frame = frame_in[conid]
geom = geom_in[conid]
body1 = geom_bodyid[geom[0]]
body2 = geom_bodyid[geom[1]]
cpos = pos_in[conid]
frame = frame_in[conid]
# TODO(team): parallelize J and Jqvel computation?
Jqvel = float(0.0)
for i in range(nv):
J = float(0.0)
jac1p, jac1r = support.jac(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
cpos,
body1,
i,
worldid,
)
jac2p, jac2r = support.jac(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
cpos,
body2,
i,
worldid,
)
for xyz in range(3):
if dimid < 3:
jac_dif = jac2p[xyz] - jac1p[xyz]
J += frame[dimid, xyz] * jac_dif
else:
jac_dif = jac2r[xyz] - jac1r[xyz]
J += frame[dimid - 3, xyz] * jac_dif
efc_J_out[worldid, efcid, i] = J
Jqvel += J * qvel_in[worldid, i]
# skip fixed bodies
body1 = body_weldid[body1]
body2 = body_weldid[body2]
da1 = body_dofadr[body1] + body_dofnum[body1] - 1
da2 = body_dofadr[body2] + body_dofnum[body2] - 1
da = wp.max(da1, da2)
for dofid in range(nv - 1, -1, -1):
if dofid == da:
# TODO(team): contact jacobian
jac1p, jac1r = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
con_pos,
body1,
dofid,
worldid,
)
jac2p, jac2r = support.jac_dof(
body_parentid,
body_rootid,
dof_bodyid,
subtree_com_in,
cdof_in,
con_pos,
body2,
dofid,
worldid,
)
J = float(0.0)
for xyz in range(3):
if dimid < 3:
jac_dif = jac2p[xyz] - jac1p[xyz]
J += frame[dimid, xyz] * jac_dif
else:
jac_dif = jac2r[xyz] - jac1r[xyz]
J += frame[dimid - 3, xyz] * jac_dif
efc_J_out[worldid, efcid, dofid] = J
Jqvel += J * qvel_in[worldid, dofid]
# Advance tree pointers and recompute da for next iteration
if da1 == da:
da1 = dof_parentid[da1]
if da2 == da:
da2 = dof_parentid[da2]
da = wp.max(da1, da2)
else:
efc_J_out[worldid, efcid, dofid] = 0.0
body_invweight0_id = worldid % body_invweight0.shape[0]
invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -1549,29 +1604,13 @@ def _efc_contact_elliptic(
)
@wp.kernel
def _num_equality(
# Data in:
ne_connect_in: wp.array(dtype=int),
ne_weld_in: wp.array(dtype=int),
ne_jnt_in: wp.array(dtype=int),
ne_ten_in: wp.array(dtype=int),
ne_flex_in: wp.array(dtype=int),
# Data out:
ne_out: wp.array(dtype=int),
):
worldid = wp.tid()
ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne_flex_in[worldid]
ne_out[worldid] = ne
@event_scope
def make_constraint(m: types.Model, d: types.Data):
"""Creates constraint jacobians and other supporting data."""
wp.launch(
_zero_constraint_counts,
dim=d.nworld,
inputs=[d.ne, d.nf, d.nl, d.nefc, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex],
inputs=[d.ne, d.nf, d.nl, d.nefc],
)
if not (m.opt.disableflags & types.DisableBit.CONSTRAINT):
@@ -1608,6 +1647,7 @@ def make_constraint(m: types.Model, d: types.Data):
refsafe,
],
outputs=[
d.ne,
d.nefc,
d.efc.type,
d.efc.id,
@@ -1618,7 +1658,6 @@ def make_constraint(m: types.Model, d: types.Data):
d.efc.vel,
d.efc.aref,
d.efc.frictionloss,
d.ne_connect,
],
)
wp.launch(
@@ -1653,6 +1692,7 @@ def make_constraint(m: types.Model, d: types.Data):
refsafe,
],
outputs=[
d.ne,
d.nefc,
d.efc.type,
d.efc.id,
@@ -1663,7 +1703,6 @@ def make_constraint(m: types.Model, d: types.Data):
d.efc.vel,
d.efc.aref,
d.efc.frictionloss,
d.ne_weld,
],
)
wp.launch(
@@ -1689,6 +1728,7 @@ def make_constraint(m: types.Model, d: types.Data):
refsafe,
],
outputs=[
d.ne,
d.nefc,
d.efc.type,
d.efc.id,
@@ -1699,7 +1739,6 @@ def make_constraint(m: types.Model, d: types.Data):
d.efc.vel,
d.efc.aref,
d.efc.frictionloss,
d.ne_jnt,
],
)
wp.launch(
@@ -1724,6 +1763,7 @@ def make_constraint(m: types.Model, d: types.Data):
refsafe,
],
outputs=[
d.ne,
d.nefc,
d.efc.type,
d.efc.id,
@@ -1734,7 +1774,6 @@ def make_constraint(m: types.Model, d: types.Data):
d.efc.vel,
d.efc.aref,
d.efc.frictionloss,
d.ne_ten,
],
)
@@ -1746,6 +1785,9 @@ def make_constraint(m: types.Model, d: types.Data):
m.opt.timestep,
m.flexedge_length0,
m.flexedge_invweight0,
m.flexedge_J_rownnz,
m.flexedge_J_rowadr,
m.flexedge_J_colind,
m.eq_solref,
m.eq_solimp,
m.eq_flex_adr,
@@ -1756,6 +1798,7 @@ def make_constraint(m: types.Model, d: types.Data):
refsafe,
],
outputs=[
d.ne,
d.nefc,
d.efc.type,
d.efc.id,
@@ -1766,17 +1809,9 @@ def make_constraint(m: types.Model, d: types.Data):
d.efc.vel,
d.efc.aref,
d.efc.frictionloss,
d.ne_flex,
],
)
wp.launch(
_num_equality,
dim=d.nworld,
inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex],
outputs=[d.ne],
)
if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS):
wp.launch(
_efc_friction_dof,
@@ -1957,8 +1992,12 @@ def make_constraint(m: types.Model, d: types.Data):
m.opt.impratio_invsqrt,
m.body_parentid,
m.body_rootid,
m.body_weldid,
m.body_dofnum,
m.body_dofadr,
m.body_invweight0,
m.dof_bodyid,
m.dof_parentid,
m.geom_bodyid,
d.qvel,
d.subtree_com,
@@ -2002,8 +2041,12 @@ def make_constraint(m: types.Model, d: types.Data):
m.opt.impratio_invsqrt,
m.body_parentid,
m.body_rootid,
m.body_weldid,
m.body_dofnum,
m.body_dofadr,
m.body_invweight0,
m.dof_bodyid,
m.dof_parentid,
m.geom_bodyid,
d.qvel,
d.subtree_com,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+12 -13
View File
@@ -25,7 +25,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -80,12 +79,12 @@ def _qderiv_actuator_passive_vel(
@cache_kernel
def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Data in:
vel_in: wp.array3d(dtype=float),
actuator_moment_in: wp.array3d(dtype=float),
# In:
vel_in: wp.array3d(dtype=float),
adr: wp.array(dtype=int),
# Out:
qDeriv_out: wp.array3d(dtype=float),
@@ -140,8 +139,8 @@ def _qderiv_actuator_passive(
# Model:
opt_timestep: wp.array(dtype=float),
opt_disableflags: int,
opt_is_sparse: bool,
dof_damping: wp.array2d(dtype=float),
is_sparse: bool,
# Data in:
qM_in: wp.array3d(dtype=float),
# In:
@@ -156,7 +155,7 @@ def _qderiv_actuator_passive(
dofiid = qMi[elemid]
dofjid = qMj[elemid]
if opt_is_sparse:
if is_sparse:
qderiv = qDeriv_in[worldid, 0, elemid]
else:
qderiv = qDeriv_in[worldid, dofiid, dofjid]
@@ -166,7 +165,7 @@ def _qderiv_actuator_passive(
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
if opt_is_sparse:
if is_sparse:
qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
else:
qM = qM_in[worldid, dofiid, dofjid] - qderiv
@@ -181,8 +180,8 @@ def _qderiv_tendon_damping(
# Model:
ntendon: int,
opt_timestep: wp.array(dtype=float),
opt_is_sparse: bool,
tendon_damping: wp.array2d(dtype=float),
is_sparse: bool,
# Data in:
ten_J_in: wp.array3d(dtype=float),
# In:
@@ -202,7 +201,7 @@ def _qderiv_tendon_damping(
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
if opt_is_sparse:
if is_sparse:
qDeriv_out[worldid, 0, elemid] -= qderiv
else:
qDeriv_out[worldid, dofiid, dofjid] -= qderiv
@@ -245,7 +244,7 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
],
outputs=[vel],
)
if m.opt.is_sparse:
if m.is_sparse:
wp.launch(
_qderiv_actuator_passive_actuation_sparse,
dim=(d.nworld, qMi.size),
@@ -258,9 +257,9 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
wp.launch_tiled(
_qderiv_actuator_passive_actuation_dense(tile, m.nu),
dim=(d.nworld, tile.adr.size),
inputs=[vel_3d, d.actuator_moment, tile.adr],
inputs=[d.actuator_moment, vel_3d, tile.adr],
outputs=[out],
block_dim=m.block_dim.mul_m_dense,
block_dim=m.block_dim.qderiv_actuator_dense,
)
wp.launch(
_qderiv_actuator_passive,
@@ -268,8 +267,8 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
inputs=[
m.opt.timestep,
m.opt.disableflags,
m.opt.is_sparse,
m.dof_damping,
m.is_sparse,
d.qM,
qMi,
qMj,
@@ -285,7 +284,7 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
wp.launch(
_qderiv_tendon_damping,
dim=(d.nworld, qMi.size),
inputs=[m.ntendon, m.opt.timestep, m.opt.is_sparse, m.tendon_damping, d.ten_J, qMi, qMj],
inputs=[m.ntendon, m.opt.timestep, m.tendon_damping, m.is_sparse, d.ten_J, qMi, qMj],
outputs=[out],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+8 -9
View File
@@ -42,7 +42,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -291,17 +290,17 @@ def _euler_damp_qfrc_sparse(
@cache_kernel
def _tile_euler_dense(tile: TileSet):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def euler_dense(
# Model:
dof_damping: wp.array2d(dtype=float),
opt_timestep: wp.array(dtype=float),
dof_damping: wp.array2d(dtype=float),
# Data in:
qM_in: wp.array3d(dtype=float),
efc_Ma_in: wp.array2d(dtype=float),
# In:
adr_in: wp.array(dtype=int),
# Out:
# Data out:
qacc_out: wp.array2d(dtype=float),
):
worldid, nodeid = wp.tid()
@@ -328,7 +327,7 @@ def euler(m: Model, d: Data):
# integrate damping implicitly
if not m.opt.disableflags & (DisableBit.EULERDAMP | DisableBit.DAMPER):
qacc = wp.empty((d.nworld, m.nv), dtype=float)
if m.opt.is_sparse:
if m.is_sparse:
qM = wp.clone(d.qM)
qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
@@ -344,7 +343,7 @@ def euler(m: Model, d: Data):
wp.launch_tiled(
_tile_euler_dense(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.dof_damping, m.opt.timestep, d.qM, d.efc.Ma, tile.adr],
inputs=[m.opt.timestep, m.dof_damping, d.qM, d.efc.Ma, tile.adr],
outputs=[qacc],
block_dim=m.block_dim.euler_dense,
)
@@ -482,7 +481,7 @@ def rungekutta4(m: Model, d: Data):
def implicit(m: Model, d: Data):
"""Integrates fully implicit in velocity."""
if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
if m.opt.is_sparse:
if m.is_sparse:
qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
else:
@@ -524,7 +523,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
# TODO(team): sparse actuator_moment version
@cache_kernel
def _actuator_velocity(nv: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def actuator_velocity(
# Data in:
qvel_in: wp.array2d(dtype=float),
@@ -544,7 +543,7 @@ def _actuator_velocity(nv: int):
@cache_kernel
def _tendon_velocity(nv: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def tendon_velocity(
# Data in:
qvel_in: wp.array2d(dtype=float),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
Vendored
+3 -5
View File
@@ -99,7 +99,7 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d(dtype=float)):
outputs=[qfrc],
)
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
if m.opt.is_sparse:
if m.is_sparse:
qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
else:
qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
@@ -120,10 +120,8 @@ def inv_constraint(m: Model, d: Data):
d.qfrc_constraint.zero_()
return
# update
h = wp.empty((d.nworld, 0, 0), dtype=float) # not used
hfactor = wp.empty((d.nworld, 0, 0), dtype=float) # not used
solver.create_context(m, d, h, hfactor, grad=False)
ctx = solver.create_inverse_context(m, d)
solver.init_context(m, d, ctx, grad=False)
def inverse(m: Model, d: Data):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+1045 -114
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+178
View File
@@ -0,0 +1,178 @@
# 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 import types
from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
@wp.kernel
def _tree_edges(
# Model:
nv: int,
body_treeid: wp.array(dtype=int),
jnt_dofadr: wp.array(dtype=int),
dof_treeid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
eq_type: wp.array(dtype=int),
eq_obj1id: wp.array(dtype=int),
eq_obj2id: wp.array(dtype=int),
eq_objtype: wp.array(dtype=int),
# Data in:
nefc_in: wp.array(dtype=int),
contact_geom_in: wp.array(dtype=wp.vec2i),
efc_type_in: wp.array2d(dtype=int),
efc_id_in: wp.array2d(dtype=int),
efc_J_in: wp.array3d(dtype=float),
njmax_in: int,
# Out:
tree_tree: wp.array3d(dtype=int), # kernel_analyzer: off
):
"""Find tree edges. Launch: (nworld, njmax)."""
worldid, efcid = wp.tid()
# skip if beyond active constraints
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
efc_type = efc_type_in[worldid, efcid]
efc_id = efc_id_in[worldid, efcid]
tree0 = int(-1)
tree1 = 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]
# site semantics
if eq_objtype[efc_id] == ObjType.SITE:
b1 = site_bodyid[b1]
b2 = site_bodyid[b2]
tree0 = body_treeid[b1]
tree1 = body_treeid[b2]
else:
# JOINT, TENDON, FLEX
use_generic = 1
# joint friction
elif efc_type == ConstraintType.FRICTION_DOF:
tree0 = dof_treeid[efc_id]
# joint limit
elif efc_type == ConstraintType.LIMIT_JOINT:
tree0 = 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]
# flex contacts have negative geom ids
if g1 >= 0 and g2 >= 0:
tree0 = body_treeid[geom_bodyid[g1]]
tree1 = body_treeid[geom_bodyid[g2]]
else:
use_generic = 1
# generic
else:
use_generic = 1
# handle static bodies
if use_generic == 0:
# swap so tree0 is non-negative if possible
if tree0 < 0 and tree1 >= 0:
tree0 = tree1
tree1 = -1
# mark the edge
if tree0 >= 0:
if tree1 < 0 or tree0 == tree1:
# self-edge
wp.atomic_max(tree_tree, worldid, tree0, tree0, 1)
else:
# cross-tree edge
t1 = wp.min(tree0, tree1)
t2 = wp.max(tree0, tree1)
wp.atomic_max(tree_tree, worldid, t1, t2, 1)
wp.atomic_max(tree_tree, worldid, t2, t1, 1)
return
# generic: scan Jacobian row
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:
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
if first_tree >= 0 and has_cross_edge == 0:
wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)):
"""Compute tree-tree adjacency matrix."""
tree_tree.zero_()
wp.launch(
kernel=_tree_edges,
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,
d.nefc,
d.contact.geom,
d.efc.type,
d.efc.id,
d.efc.J,
d.njmax,
],
outputs=[tree_tree],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+6 -6
View File
@@ -258,7 +258,7 @@ def _gravity_force(
if gravcomp:
force = -gravity * body_mass[worldid % body_mass.shape[0], bodyid] * gravcomp
pos = xipos_in[worldid, bodyid]
jac, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos, bodyid, dofid, worldid)
jac, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos, bodyid, dofid, worldid)
wp.atomic_add(qfrc_gravcomp_out[worldid], dofid, wp.dot(jac, force))
@@ -526,9 +526,9 @@ def _fluid(m: Model, d: Data):
@wp.kernel
def _qfrc_passive(
# Model:
opt_has_fluid: bool,
jnt_actgravcomp: wp.array(dtype=int),
dof_jntid: wp.array(dtype=int),
has_fluid: bool,
# Data in:
qfrc_spring_in: wp.array2d(dtype=float),
qfrc_damper_in: wp.array2d(dtype=float),
@@ -548,7 +548,7 @@ def _qfrc_passive(
qfrc_passive += qfrc_gravcomp_in[worldid, dofid]
# add fluid force
if opt_has_fluid:
if has_fluid:
qfrc_passive += qfrc_fluid_in[worldid, dofid]
qfrc_passive_out[worldid, dofid] = qfrc_passive
@@ -703,7 +703,7 @@ def _flex_bending(
force[i, x] -= flex_bending[edgeid, 16] * frc[i, x]
for i in range(nvert):
bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]
bodyid = flex_vertbodyid[v[i]]
for x in range(3):
wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[i, x])
@@ -826,16 +826,16 @@ def passive(m: Model, d: Data):
outputs=[d.qfrc_gravcomp],
)
if m.opt.has_fluid:
if m.has_fluid:
_fluid(m, d)
wp.launch(
_qfrc_passive,
dim=(d.nworld, m.nv),
inputs=[
m.opt.has_fluid,
m.jnt_actgravcomp,
m.dof_jntid,
m.has_fluid,
d.qfrc_spring,
d.qfrc_damper,
d.qfrc_gravcomp,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+352 -68
View File
@@ -13,15 +13,17 @@
# limitations under the License.
# ==============================================================================
from typing import Optional, Tuple
from typing import Tuple
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_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 GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
wp.set_module_options({"enable_backward": False})
@@ -183,7 +185,7 @@ def _ray_triangle(
@wp.func
def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
def ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance and normal at which a ray intersects with a plane."""
# map to local frame
lpnt, lvec = _ray_map(pos, mat, pnt, vec)
@@ -207,7 +209,7 @@ def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp
@wp.func
def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
def ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance and normal at which a ray intersects with a sphere."""
dif = pnt - pos
@@ -224,11 +226,11 @@ def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tu
@wp.func
def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
def ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance and normal at which a ray intersects with a capsule."""
# bounding sphere test
ssz = size[0] + size[1]
dist_sphere, normal_sphere = _ray_sphere(pos, ssz * ssz, pnt, vec)
dist_sphere, normal_sphere = ray_sphere(pos, ssz * ssz, pnt, vec)
if dist_sphere < 0:
return -1.0, wp.vec3()
@@ -248,8 +250,9 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
sol, xx = _ray_quad(a, b, c)
part = 0 # -1: bottom, 0: cylinder, 1: top
# make sure round solution is between flat sides
if sol >= 0.0 and wp.abs(lpnt[2] + sol * vec[2]) <= size[1]:
# make sure round solution is between flat sides (must use local z component)
# TODO: We should add a test to catch this case.
if sol >= 0.0 and wp.abs(lpnt[2] + sol * lvec[2]) <= size[1]:
if x < 0.0 or sol < x:
x = sol
@@ -260,7 +263,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
c = wp.dot(ldif, ldif) - sq_size0
_, xx = _ray_quad(a, b, c)
# accept only top half of sphere
# accept only top half of sphere (use local z component)
for i in range(2):
if xx[i] >= 0.0 and lpnt[2] + xx[i] * lvec[2] >= size[1]:
if x < 0.0 or xx[i] < x:
@@ -273,7 +276,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
c = wp.dot(ldif, ldif) - sq_size0
_, xx = _ray_quad(a, b, c)
# accept only bottom half of sphere
# accept only bottom half of sphere (use local z component)
for i in range(2):
if xx[i] >= 0.0 and lpnt[2] + xx[i] * lvec[2] <= -size[1]:
if x < 0.0 or xx[i] < x:
@@ -297,7 +300,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
@wp.func
def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
def ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance and normal at which a ray intersects with an ellipsoid."""
# map to local frame
lpnt, lvec = _ray_map(pos, mat, pnt, vec)
@@ -328,11 +331,11 @@ def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec
@wp.func
def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
def ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance and normal at which a ray intersects with a cylinder."""
# bounding sphere test
ssz = size[0] * size[0] + size[1] * size[1]
dist_sphere, normal_sphere = _ray_sphere(pos, ssz, pnt, vec)
dist_sphere, normal_sphere = ray_sphere(pos, ssz, pnt, vec)
if dist_sphere < 0:
return -1.0, wp.vec3()
@@ -392,13 +395,13 @@ _IFACE = wp.types.matrix((3, 2), dtype=int)(1, 2, 0, 2, 0, 1)
@wp.func
def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
def ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
"""Returns distance, per side information, and normal at which a ray intersects with a box."""
all = vec6(-1.0, -1.0, -1.0, -1.0, -1.0, -1.0)
# bounding sphere test
ssz = wp.dot(size, size)
dist_sphere, _ = _ray_sphere(pos, ssz, pnt, vec)
dist_sphere, _ = ray_sphere(pos, ssz, pnt, vec)
if dist_sphere < 0:
return -1.0, all, wp.vec3()
@@ -446,7 +449,7 @@ def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
@wp.func
def _ray_hfield(
def ray_hfield(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
@@ -487,10 +490,10 @@ def _ray_hfield(
top_pos = pos + mat_col * top_scale
# init: intersection with base box
x, _, normal_base = _ray_box(base_pos, mat, base_size, pnt, vec)
x, _, normal_base = ray_box(base_pos, mat, base_size, pnt, vec)
# check top box: done if no intersection
top_intersect, all, normal_top = _ray_box(top_pos, mat, top_size, pnt, vec)
top_intersect, all, normal_top = ray_box(top_pos, mat, top_size, pnt, vec)
if top_intersect < 0.0:
return x, normal_base
@@ -627,10 +630,16 @@ def ray_mesh(
data_id: int,
pos: wp.vec3,
mat: wp.mat33,
size: wp.vec3,
pnt: wp.vec3,
vec: wp.vec3,
) -> Tuple[float, wp.vec3]:
"""Returns the distance and normal for ray mesh intersections."""
# bounding box test
dist_box, _all, _normal = ray_box(pos, mat, size, pnt, vec)
if dist_box < 0.0:
return -1.0, wp.vec3()
pnt, vec = _ray_map(pos, mat, pnt, vec)
# compute orthogonal basis vectors
@@ -687,6 +696,87 @@ def ray_mesh(
return x, normal
@wp.func
def ray_mesh_with_bvh(
# In:
mesh_bvh_id: wp.array(dtype=wp.uint64),
mesh_geom_id: int,
pos: wp.vec3,
mat: wp.mat33,
pnt: wp.vec3,
vec: wp.vec3,
max_t: float,
) -> 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.
"""
t = float(-1.0)
u = float(0.0)
v = float(0.0)
sign = float(0.0)
n = wp.vec3(0.0, 0.0, 0.0)
f = int(-1)
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
normal = mat @ n
normal = wp.normalize(normal)
return t, normal, u, v, f, mesh_geom_id
return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
@wp.func
def ray_mesh_with_bvh_anyhit(
# In:
mesh_bvh_id: wp.array(dtype=wp.uint64),
mesh_geom_id: int,
pos: wp.vec3,
mat: wp.mat33,
pnt: wp.vec3,
vec: wp.vec3,
max_t: float,
) -> bool:
"""Returns True if there is any hit for ray mesh intersections.
Requires wp.Mesh be constructed and their ids to be passed. This variant is useful
for shadow ray casts where the only goal is if there is any ray hit.
"""
lpnt, lvec = _ray_map(pos, mat, pnt, vec)
return wp.mesh_query_ray_anyhit(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t)
@wp.func
def ray_flex_with_bvh(
# In:
bvh_id: wp.uint64,
group_root: int,
pnt: wp.vec3,
vec: wp.vec3,
max_t: float,
) -> Tuple[float, wp.vec3, float, float, int]:
"""Returns intersection information for flex intersections.
Requires wp.Mesh be constructed and their ids to be passed. Flex are already in world space.
"""
t = float(-1.0)
u = float(0.0)
v = float(0.0)
sign = float(0.0)
n = wp.vec3(0.0, 0.0, 0.0)
f = int(-1)
hit = wp.mesh_query_ray(bvh_id, pnt, vec, max_t, t, u, v, sign, n, f, group_root)
if hit:
return t, n, u, v, f
return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1
@wp.func
def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> Tuple[float, wp.vec3]:
"""Returns distance along ray to intersection with geom and normal at intersection point.
@@ -695,17 +785,17 @@ def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
"""
# TODO(team): static loop unrolling to remove unnecessary branching
if geomtype == GeomType.PLANE:
return _ray_plane(pos, mat, size, pnt, vec)
return ray_plane(pos, mat, size, pnt, vec)
elif geomtype == GeomType.SPHERE:
return _ray_sphere(pos, size[0] * size[0], pnt, vec)
return ray_sphere(pos, size[0] * size[0], pnt, vec)
elif geomtype == GeomType.CAPSULE:
return _ray_capsule(pos, mat, size, pnt, vec)
return ray_capsule(pos, mat, size, pnt, vec)
elif geomtype == GeomType.ELLIPSOID:
return _ray_ellipsoid(pos, mat, size, pnt, vec)
return ray_ellipsoid(pos, mat, size, pnt, vec)
elif geomtype == GeomType.CYLINDER:
return _ray_cylinder(pos, mat, size, pnt, vec)
return ray_cylinder(pos, mat, size, pnt, vec)
elif geomtype == GeomType.BOX:
dist, _, normal = _ray_box(pos, mat, size, pnt, vec)
dist, _, normal = ray_box(pos, mat, size, pnt, vec)
return dist, normal
else:
return -1.0, wp.vec3()
@@ -771,11 +861,12 @@ def _ray_geom_mesh(
geom_dataid[geomid],
pos,
mat,
geom_size[worldid % geom_size.shape[0], geomid],
pnt,
vec,
)
elif type == GeomType.HFIELD:
return _ray_hfield(
return ray_hfield(
geom_type,
geom_dataid,
hfield_size,
@@ -836,7 +927,7 @@ def _ray(
num_threads = wp.block_dim()
min_dist = float(wp.inf)
min_dist = float(MJ_MAXVAL)
min_geomid = int(-1)
min_normal = wp.vec3()
@@ -874,9 +965,9 @@ def _ray(
geomid,
)
if dist < 0:
dist = wp.inf
dist = MJ_MAXVAL
else:
dist = wp.inf
dist = MJ_MAXVAL
normal = wp.vec3()
tile_dist = wp.tile(dist)
@@ -891,7 +982,162 @@ def _ray(
min_geomid = tile_geomid[local_min_geomid[0]]
min_normal = tile_normal[local_min_geomid[0]]
if wp.isinf(min_dist):
if min_dist >= MJ_MAXVAL:
dist_out[worldid, rayid] = -1.0
else:
dist_out[worldid, rayid] = min_dist
geomid_out[worldid, rayid] = min_geomid
normal_out[worldid, rayid] = min_normal
@wp.func
def _ray_geom_mesh_bvh(
# Model:
body_weldid: wp.array(dtype=int),
geom_type: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_matid: wp.array2d(dtype=int),
geom_group: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_rgba: wp.array2d(dtype=wp.vec4),
mat_rgba: wp.array2d(dtype=wp.vec4),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
worldid: int,
pnt: wp.vec3,
vec: wp.vec3,
geomgroup: vec6,
flg_static: bool,
bodyexclude: int,
geomid: int,
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
min_dist: float,
) -> Tuple[float, wp.vec3]:
if not _ray_eliminate(
body_weldid,
geom_bodyid,
geom_matid[worldid % geom_matid.shape[0]],
geom_group,
geom_rgba[worldid % geom_rgba.shape[0]],
mat_rgba[worldid % mat_rgba.shape[0]],
geomid,
geomgroup,
flg_static,
bodyexclude,
):
pos = geom_xpos_in[worldid, geomid]
mat = geom_xmat_in[worldid, geomid]
gtype = geom_type[geomid]
if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
bvh_ids,
geom_dataid[geomid],
pos,
mat,
pnt,
vec,
min_dist,
)
if t >= 0.0 and t < min_dist:
return t, n
else:
return ray_geom(
pos,
mat,
geom_size[worldid % geom_size.shape[0], geomid],
pnt,
vec,
gtype,
)
return -1.0, wp.vec3()
@wp.kernel
def _ray_bvh(
# Model:
ngeom: int,
body_weldid: wp.array(dtype=int),
geom_type: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_matid: wp.array2d(dtype=int),
geom_group: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_rgba: wp.array2d(dtype=wp.vec4),
mat_rgba: wp.array2d(dtype=wp.vec4),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
pnt: wp.array2d(dtype=wp.vec3),
vec: wp.array2d(dtype=wp.vec3),
geomgroup: vec6,
flg_static: bool,
bodyexclude: wp.array(dtype=int),
bvh_id: wp.uint64,
group_root: wp.array(dtype=int),
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
# Out:
dist_out: wp.array2d(dtype=float),
geomid_out: wp.array2d(dtype=int),
normal_out: wp.array2d(dtype=wp.vec3),
):
worldid, rayid = wp.tid()
ray_origin = pnt[worldid, rayid]
ray_dir = vec[worldid, rayid]
body_exclude = bodyexclude[rayid]
min_dist = float(MJ_MAXVAL)
min_geomid = int(-1)
min_normal = wp.vec3()
query = wp.bvh_query_ray(bvh_id, ray_origin, ray_dir, group_root[worldid])
bounds_nr = int(0)
while wp.bvh_query_next(query, bounds_nr, min_dist):
bvh_local = bounds_nr - (worldid * ngeom)
geomid = enabled_geom_ids[bvh_local]
dist, normal = _ray_geom_mesh_bvh(
body_weldid,
geom_type,
geom_bodyid,
geom_dataid,
geom_matid,
geom_group,
geom_size,
geom_rgba,
mat_rgba,
geom_xpos_in,
geom_xmat_in,
worldid,
pnt[worldid, rayid],
vec[worldid, rayid],
geomgroup,
flg_static,
body_exclude,
geomid,
mesh_bvh_id,
hfield_bvh_id,
min_dist,
)
if dist >= 0.0 and dist < min_dist:
min_dist = dist
min_geomid = geomid
min_normal = normal
if min_dist >= MJ_MAXVAL:
dist_out[worldid, rayid] = -1.0
else:
dist_out[worldid, rayid] = min_dist
@@ -904,9 +1150,10 @@ def ray(
d: Data,
pnt: wp.array2d(dtype=wp.vec3),
vec: wp.array2d(dtype=wp.vec3),
geomgroup: Optional[vec6] = None,
geomgroup: vec6 | None = None,
flg_static: bool = True,
bodyexclude: int = -1,
rc: RenderContext | None = None,
) -> Tuple[wp.array, wp.array, wp.array]:
"""Returns the distance at which rays intersect with primitive geoms.
@@ -915,9 +1162,11 @@ def ray(
d: The data object containing the current state and output arrays (device).
pnt: Ray origin points.
vec: Ray directions.
geomgroup: Group inclusion/exclusion mask. If all are wp.inf, ignore.
geomgroup: Group inclusion/exclusion mask.
flg_static: If True, allows rays to intersect with static geoms.
bodyexclude: Ignore geoms on specified body id (-1 to disable).
rc: Optional Render context containing BVH information for BVH accelerated ray
intersections.
Returns:
Distances from ray origins to geom surfaces, IDs of intersected geoms (-1 if none),
@@ -935,7 +1184,7 @@ def ray(
ray_geomid = wp.empty((d.nworld, 1), dtype=int)
ray_normal = wp.empty((d.nworld, 1), dtype=wp.vec3)
rays(m, d, pnt, vec, geomgroup, flg_static, ray_bodyexclude, ray_dist, ray_geomid, ray_normal)
rays(m, d, pnt, vec, geomgroup, flg_static, ray_bodyexclude, ray_dist, ray_geomid, ray_normal, rc)
return ray_dist, ray_geomid, ray_normal
@@ -951,6 +1200,7 @@ def rays(
dist: wp.array2d(dtype=float),
geomid: wp.array2d(dtype=int),
normal: wp.array2d(dtype=wp.vec3),
rc: RenderContext | None = None,
):
"""Ray intersection for multiple worlds and multiple rays.
@@ -968,41 +1218,75 @@ def rays(
geomid: Output array for IDs of intersected geoms, shape (nworld, nray). -1
indicates no intersection.
normal: Output array for normals at intersection points, shape (nworld, nray).
rc: Optional Render context containing BVH information for BVH accelerated ray
intersections.
"""
wp.launch_tiled(
_ray,
dim=(d.nworld, pnt.shape[1]),
inputs=[
m.ngeom,
m.nmeshface,
m.body_weldid,
m.geom_type,
m.geom_bodyid,
m.geom_dataid,
m.geom_matid,
m.geom_group,
m.geom_size,
m.geom_rgba,
m.mesh_vertadr,
m.mesh_faceadr,
m.mesh_vert,
m.mesh_face,
m.hfield_size,
m.hfield_nrow,
m.hfield_ncol,
m.hfield_adr,
m.hfield_data,
m.mat_rgba,
d.geom_xpos,
d.geom_xmat,
pnt,
vec,
geomgroup,
flg_static,
bodyexclude,
dist,
geomid,
normal,
],
block_dim=m.block_dim.ray,
)
# TODO: Investigate building rc if none and removing the non-accelerated path
if rc is None:
wp.launch_tiled(
_ray,
dim=(d.nworld, pnt.shape[1]),
inputs=[
m.ngeom,
m.nmeshface,
m.body_weldid,
m.geom_type,
m.geom_bodyid,
m.geom_dataid,
m.geom_matid,
m.geom_group,
m.geom_size,
m.geom_rgba,
m.mesh_vertadr,
m.mesh_faceadr,
m.mesh_vert,
m.mesh_face,
m.hfield_size,
m.hfield_nrow,
m.hfield_ncol,
m.hfield_adr,
m.hfield_data,
m.mat_rgba,
d.geom_xpos,
d.geom_xmat,
pnt,
vec,
geomgroup,
flg_static,
bodyexclude,
dist,
geomid,
normal,
],
block_dim=m.block_dim.ray,
)
else:
wp.launch(
_ray_bvh,
dim=(d.nworld, pnt.shape[1]),
inputs=[
rc.bvh_ngeom,
m.body_weldid,
m.geom_type,
m.geom_bodyid,
m.geom_dataid,
m.geom_matid,
m.geom_group,
m.geom_size,
m.geom_rgba,
m.mat_rgba,
d.geom_xpos,
d.geom_xmat,
pnt,
vec,
geomgroup,
flg_static,
bodyexclude,
rc.bvh_id,
rc.group_root,
rc.enabled_geom_ids,
rc.mesh_bvh_id,
rc.hfield_bvh_id,
],
outputs=[dist, geomid, normal],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+696
View File
@@ -0,0 +1,696 @@
# 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.
# ==============================================================================
from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_box
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_capsule
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_cylinder
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_ellipsoid
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_flex_with_bvh
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh_with_bvh
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh_with_bvh_anyhit
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_plane
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_sphere
from mujoco.mjx.third_party.mujoco_warp._src.render_util import compute_ray
from mujoco.mjx.third_party.mujoco_warp._src.render_util import pack_rgba_to_uint32
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
from mujoco.mjx.third_party.mujoco_warp._src.types import TEXTURE_DTYPE
@wp.func
def sample_texture(
# Model:
geom_type: wp.array(dtype=int),
mesh_faceadr: wp.array(dtype=int),
# In:
geom_id: int,
tex_repeat: wp.vec2,
tex: TEXTURE_DTYPE,
pos: wp.vec3,
rot: wp.mat33,
mesh_facetexcoord: wp.array(dtype=wp.vec3i),
mesh_texcoord: wp.array(dtype=wp.vec2),
mesh_texcoord_offsets: wp.array(dtype=int),
hit_point: wp.vec3,
bary_u: float,
bary_v: float,
f: int,
mesh_id: int,
) -> wp.vec3:
uv = wp.vec2(0.0, 0.0)
if geom_type[geom_id] == GeomType.PLANE:
local = wp.transpose(rot) @ (hit_point - pos)
uv = wp.vec2(local[0], local[1])
if geom_type[geom_id] == GeomType.MESH:
if f < 0 or mesh_id < 0:
return wp.vec3(0.0, 0.0, 0.0)
face_adr = mesh_faceadr[mesh_id] + f
uv0 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][0]]
uv1 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][1]]
uv2 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][2]]
uv = uv0 * bary_u + uv1 * bary_v + uv2 * (1.0 - bary_u - bary_v)
u = uv[0] * tex_repeat[0]
v = uv[1] * tex_repeat[1]
u = u - wp.floor(u)
v = v - wp.floor(v)
tex_color = wp.texture_sample(tex, wp.vec2(u, v), dtype=wp.vec4)
return wp.vec3(tex_color[0], tex_color[1], tex_color[2])
# TODO: Investigate combining cast_ray and cast_ray_first_hit
@wp.func
def cast_ray(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
bvh_id: wp.uint64,
group_root: int,
world_id: int,
bvh_ngeom: int,
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
) -> Tuple[int, float, wp.vec3, float, float, int, int]:
dist = float(MJ_MAXVAL)
normal = wp.vec3(0.0, 0.0, 0.0)
geom_id = int(-1)
bary_u = float(0.0)
bary_v = float(0.0)
face_idx = int(-1)
geom_mesh_id = int(-1)
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
bounds_nr = int(0)
while wp.bvh_query_next(query, bounds_nr, dist):
gi_global = bounds_nr
gi_bvh_local = gi_global - (world_id * bvh_ngeom)
gi = enabled_geom_ids[gi_bvh_local]
hit_mesh_id = int(-1)
u = float(0.0)
v = float(0.0)
f = int(-1)
n = wp.vec3(0.0, 0.0, 0.0)
# TODO: Investigate branch elimination with static loop unrolling
if geom_type[gi] == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
ray_origin_world,
ray_dir_world,
dist,
)
if geom_type[gi] == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.MESH:
d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
mesh_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
ray_origin_world,
ray_dir_world,
dist,
)
if d >= 0.0 and d < dist:
dist = d
normal = n
geom_id = gi
bary_u = u
bary_v = v
face_idx = f
geom_mesh_id = hit_mesh_id
return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id
@wp.func
def cast_ray_first_hit(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
bvh_id: wp.uint64,
group_root: int,
world_id: int,
bvh_ngeom: int,
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
max_dist: float,
) -> 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)
bounds_nr = int(0)
while wp.bvh_query_next(query, bounds_nr, max_dist):
gi_global = bounds_nr
gi_bvh_local = gi_global - (world_id * bvh_ngeom)
gi = enabled_geom_ids[gi_bvh_local]
# TODO: Investigate branch elimination with static loop unrolling
if geom_type[gi] == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
ray_origin_world,
ray_dir_world,
max_dist,
)
if geom_type[gi] == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
geom_size[world_id % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if geom_type[gi] == GeomType.MESH:
hit = ray_mesh_with_bvh_anyhit(
mesh_bvh_id,
geom_dataid[gi],
geom_xpos_in[world_id, gi],
geom_xmat_in[world_id, gi],
ray_origin_world,
ray_dir_world,
max_dist,
)
d = 0.0 if hit else -1.0
if d >= 0.0 and d < max_dist:
return True
return False
@wp.func
def compute_lighting(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
use_shadows: bool,
bvh_id: wp.uint64,
group_root: int,
bvh_ngeom: int,
enabled_geom_ids: wp.array(dtype=int),
world_id: int,
mesh_bvh_id: wp.array(dtype=wp.uint64),
hfield_bvh_id: wp.array(dtype=wp.uint64),
lightactive: bool,
lighttype: int,
lightcastshadow: bool,
lightpos: wp.vec3,
lightdir: wp.vec3,
normal: wp.vec3,
hitpoint: wp.vec3,
) -> float:
light_contribution = float(0.0)
# TODO: We should probably only be looping over active lights
# in the first place with a static loop of enabled light idx?
if not lightactive:
return light_contribution
L = wp.vec3(0.0, 0.0, 0.0)
dist_to_light = float(MJ_MAXVAL)
attenuation = float(1.0)
if lighttype == 1: # directional light
L = wp.normalize(-lightdir)
else:
L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint)
attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light)
if lighttype == 0: # spot light
spot_dir = wp.normalize(lightdir)
cos_theta = wp.dot(-L, spot_dir)
spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) / (0.95 - 0.85)))
attenuation = attenuation * spot_factor
ndotl = wp.max(0.0, wp.dot(normal, L))
if ndotl == 0.0:
return light_contribution
visible = float(1.0)
if use_shadows and lightcastshadow:
# Nudge the origin slightly along the surface normal to avoid
# self-intersection when casting shadow rays
eps = 1.0e-4
shadow_origin = hitpoint + normal * eps
# Distance-limited shadows: cap by dist_to_light (for non-directional)
max_t = float(dist_to_light - 1.0e-3)
if lighttype == 1: # directional light
max_t = float(1.0e8)
shadow_hit = cast_ray_first_hit(
geom_type,
geom_dataid,
geom_size,
geom_xpos_in,
geom_xmat_in,
bvh_id,
group_root,
world_id,
bvh_ngeom,
enabled_geom_ids,
mesh_bvh_id,
hfield_bvh_id,
shadow_origin,
L,
max_t,
)
if shadow_hit:
visible = 0.3
return ndotl * attenuation * visible
@event_scope
def render(m: Model, d: Data, rc: RenderContext):
"""Render the current frame.
Outputs are stored in buffers within the render context.
Args:
m: The model on device.
d: The data on device.
rc: The render context on device.
"""
rc.rgb_data.fill_(rc.background_color)
rc.depth_data.fill_(0.0)
@wp.kernel(module="unique", enable_backward=False)
def _render_megakernel(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_matid: wp.array2d(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_rgba: wp.array2d(dtype=wp.vec4),
cam_projection: wp.array(dtype=int),
cam_fovy: wp.array2d(dtype=float),
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_intrinsic: wp.array2d(dtype=wp.vec4),
light_type: wp.array2d(dtype=int),
light_castshadow: wp.array2d(dtype=bool),
light_active: wp.array2d(dtype=bool),
mesh_faceadr: wp.array(dtype=int),
mat_texid: wp.array3d(dtype=int),
mat_texrepeat: wp.array2d(dtype=wp.vec2),
mat_rgba: wp.array2d(dtype=wp.vec4),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
light_xpos_in: wp.array2d(dtype=wp.vec3),
light_xdir_in: wp.array2d(dtype=wp.vec3),
# In:
nrender: int,
use_shadows: bool,
bvh_ngeom: int,
cam_res: wp.array(dtype=wp.vec2i),
cam_id_map: wp.array(dtype=int),
ray: wp.array(dtype=wp.vec3),
rgb_adr: wp.array(dtype=int),
depth_adr: wp.array(dtype=int),
render_rgb: wp.array(dtype=bool),
render_depth: wp.array(dtype=bool),
bvh_id: wp.uint64,
group_root: wp.array(dtype=int),
flex_bvh_id: wp.uint64,
flex_group_root: wp.array(dtype=int),
enabled_geom_ids: wp.array(dtype=int),
mesh_bvh_id: wp.array(dtype=wp.uint64),
mesh_facetexcoord: wp.array(dtype=wp.vec3i),
mesh_texcoord: wp.array(dtype=wp.vec2),
mesh_texcoord_offsets: wp.array(dtype=int),
hfield_bvh_id: wp.array(dtype=wp.uint64),
flex_rgba: wp.array(dtype=wp.vec4),
# TODO: remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
textures: wp.array(dtype=TEXTURE_DTYPE),
# Out:
rgb_out: wp.array2d(dtype=wp.uint32),
depth_out: wp.array2d(dtype=float),
):
world_idx, ray_idx = wp.tid()
# Map global ray_idx -> (cam_idx, ray_idx_local) using cumulative sizes
cam_idx = int(-1)
ray_idx_local = int(-1)
accum = int(0)
for i in range(nrender):
num_i = cam_res[i][0] * cam_res[i][1]
if ray_idx < accum + num_i:
cam_idx = i
ray_idx_local = ray_idx - accum
break
accum += num_i
if cam_idx == -1 or ray_idx_local < 0:
return
if not render_rgb[cam_idx] and not render_depth[cam_idx]:
return
# Map active camera index to MuJoCo camera ID
mujoco_cam_id = cam_id_map[cam_idx]
if wp.static(rc.ray is None):
img_w = cam_res[cam_idx][0]
img_h = cam_res[cam_idx][1]
px = ray_idx_local % img_w
py = ray_idx_local // img_w
ray_dir_local_cam = compute_ray(
cam_projection[mujoco_cam_id],
cam_fovy[world_idx % cam_fovy.shape[0], mujoco_cam_id],
cam_sensorsize[mujoco_cam_id],
cam_intrinsic[world_idx % cam_intrinsic.shape[0], mujoco_cam_id],
img_w,
img_h,
px,
py,
wp.static(rc.znear),
)
else:
ray_dir_local_cam = ray[ray_idx]
ray_dir_world = cam_xmat_in[world_idx, mujoco_cam_id] @ ray_dir_local_cam
ray_origin_world = cam_xpos_in[world_idx, mujoco_cam_id]
geom_id, dist, normal, u, v, f, mesh_id = cast_ray(
geom_type,
geom_dataid,
geom_size,
geom_xpos_in,
geom_xmat_in,
bvh_id,
group_root[world_idx],
world_idx,
bvh_ngeom,
enabled_geom_ids,
mesh_bvh_id,
hfield_bvh_id,
ray_origin_world,
ray_dir_world,
)
if wp.static(m.nflex > 0):
d, n, u, v, f = ray_flex_with_bvh(
flex_bvh_id,
flex_group_root[world_idx],
ray_origin_world,
ray_dir_world,
dist,
)
if d >= 0.0 and d < dist:
dist = d
normal = n
geom_id = -2
# Early Out
if geom_id == -1:
return
if render_depth[cam_idx]:
depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist
if not render_rgb[cam_idx]:
return
# Shade the pixel
hit_point = ray_origin_world + ray_dir_world * dist
if geom_id == -2:
# TODO: Currently flex textures are not supported, and only the first rgba value
# is used until further flex support is added.
color = flex_rgba[0]
elif geom_matid[world_idx % geom_matid.shape[0], geom_id] == -1:
color = geom_rgba[world_idx % geom_rgba.shape[0], geom_id]
else:
color = mat_rgba[world_idx % mat_rgba.shape[0], geom_matid[world_idx % geom_matid.shape[0], geom_id]]
base_color = wp.vec3(color[0], color[1], color[2])
hit_color = base_color
if wp.static(rc.use_textures):
if geom_id != -2:
mat_id = geom_matid[world_idx % geom_matid.shape[0], geom_id]
if mat_id >= 0:
tex_id = mat_texid[world_idx % mat_texid.shape[0], mat_id, 1]
if tex_id >= 0:
tex_color = sample_texture(
geom_type,
mesh_faceadr,
geom_id,
mat_texrepeat[world_idx % mat_texrepeat.shape[0], mat_id],
textures[tex_id],
geom_xpos_in[world_idx, geom_id],
geom_xmat_in[world_idx, geom_id],
mesh_facetexcoord,
mesh_texcoord,
mesh_texcoord_offsets,
hit_point,
u,
v,
f,
mesh_id,
)
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)
# Apply lighting and shadows
for l in range(wp.static(m.nlight)):
light_contribution = compute_lighting(
geom_type,
geom_dataid,
geom_size,
geom_xpos_in,
geom_xmat_in,
use_shadows,
bvh_id,
group_root[world_idx],
bvh_ngeom,
enabled_geom_ids,
world_idx,
mesh_bvh_id,
hfield_bvh_id,
light_active[world_idx % light_active.shape[0], l],
light_type[world_idx % light_type.shape[0], l],
light_castshadow[world_idx % light_castshadow.shape[0], l],
light_xpos_in[world_idx, l],
light_xdir_in[world_idx, l],
normal,
hit_point,
)
result = result + base_color * light_contribution
hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0))
hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0))
rgb_out[world_idx, rgb_adr[cam_idx] + ray_idx_local] = pack_rgba_to_uint32(
hit_color[0] * 255.0,
hit_color[1] * 255.0,
hit_color[2] * 255.0,
255.0,
)
wp.launch(
kernel=_render_megakernel,
dim=(d.nworld, rc.total_rays),
inputs=[
m.geom_type,
m.geom_dataid,
m.geom_matid,
m.geom_size,
m.geom_rgba,
m.cam_projection,
m.cam_fovy,
m.cam_sensorsize,
m.cam_intrinsic,
m.light_type,
m.light_castshadow,
m.light_active,
m.mesh_faceadr,
m.mat_texid,
m.mat_texrepeat,
m.mat_rgba,
d.geom_xpos,
d.geom_xmat,
d.cam_xpos,
d.cam_xmat,
d.light_xpos,
d.light_xdir,
rc.nrender,
rc.use_shadows,
rc.bvh_ngeom,
rc.cam_res,
rc.cam_id_map,
rc.ray,
rc.rgb_adr,
rc.depth_adr,
rc.render_rgb,
rc.render_depth,
rc.bvh_id,
rc.group_root,
rc.flex_bvh_id,
rc.flex_group_root,
rc.enabled_geom_ids,
rc.mesh_bvh_id,
rc.mesh_facetexcoord,
rc.mesh_texcoord,
rc.mesh_texcoord_offsets,
rc.hfield_bvh_id,
rc.flex_rgba,
rc.textures,
],
outputs=[
rc.rgb_data,
rc.depth_data,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
Vendored
+130
View File
@@ -0,0 +1,130 @@
# 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 mujoco
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.types import ProjectionType
wp.set_module_options({"enable_backward": False})
@wp.kernel
def _convert_texture_data(
# In:
width: int,
adr: int,
nc: int,
tex_data_in: wp.array(dtype=wp.uint8),
# Out:
tex_data_out: wp.array3d(dtype=float),
):
"""Convert uint8 texture data to vec4 format for efficient sampling."""
x, y = wp.tid()
offset = adr + (y * width + x) * nc
r = tex_data_in[offset + 0] if nc > 0 else wp.uint8(0)
g = tex_data_in[offset + 1] if nc > 1 else wp.uint8(0)
b = tex_data_in[offset + 2] if nc > 2 else wp.uint8(0)
a = wp.uint8(255)
tex_data_out[y, x, 0] = float(r) * wp.static(1.0 / 255.0)
tex_data_out[y, x, 1] = float(g) * wp.static(1.0 / 255.0)
tex_data_out[y, x, 2] = float(b) * wp.static(1.0 / 255.0)
tex_data_out[y, x, 3] = float(a) * wp.static(1.0 / 255.0)
def create_warp_texture(mjm: mujoco.MjModel, tex_id: int) -> wp.array:
"""Create a Warp texture from a MuJoCo model texture data."""
tex_adr = mjm.tex_adr[tex_id]
tex_width = mjm.tex_width[tex_id]
tex_height = mjm.tex_height[tex_id]
nchannel = mjm.tex_nchannel[tex_id]
tex_data = wp.zeros((tex_height, tex_width, 4), dtype=float)
wp.launch(
_convert_texture_data,
dim=(tex_width, tex_height),
inputs=[tex_width, tex_adr, nchannel, wp.array(mjm.tex_data, dtype=wp.uint8)],
outputs=[tex_data],
)
return wp.Texture2D(tex_data, filter_mode=wp.TextureFilterMode.LINEAR)
@wp.func
def compute_ray(
# In:
projection: int,
fovy: float,
sensorsize: wp.vec2,
intrinsic: wp.vec4,
img_w: int,
img_h: int,
px: int,
py: int,
znear: float,
) -> wp.vec3:
"""Compute ray direction for a pixel with per-world camera parameters.
This combines _camera_frustum_bounds and build_primary_rays logic for use
inside a kernel when camera parameters are batched/randomized across worlds.
"""
if projection == ProjectionType.ORTHOGRAPHIC:
return wp.vec3(0.0, 0.0, -1.0)
aspect = float(img_w) / float(img_h)
sensor_h = sensorsize[1]
# Check if we have intrinsics (sensorsize[1] != 0)
if sensor_h != 0.0:
fx = intrinsic[0]
fy = intrinsic[1]
cx = intrinsic[2]
cy = intrinsic[3]
sensor_w = sensorsize[0]
target_aspect = float(img_w) / float(img_h)
sensor_aspect = sensor_w / sensor_h
if target_aspect > sensor_aspect:
sensor_h = sensor_w / target_aspect
elif target_aspect < sensor_aspect:
sensor_w = sensor_h * target_aspect
inv_fx_znear = znear / fx
inv_fy_znear = znear / fy
left = -inv_fx_znear * (sensor_w * 0.5 - cx)
right = inv_fx_znear * (sensor_w * 0.5 + cx)
top = inv_fy_znear * (sensor_h * 0.5 - cy)
bottom = -inv_fy_znear * (sensor_h * 0.5 + cy)
else:
fovy_rad = fovy * wp.static(wp.pi / 180.0)
half_height = znear * wp.tan(0.5 * fovy_rad)
half_width = half_height * aspect
left = -half_width
right = half_width
top = half_height
bottom = -half_height
u = (float(px) + 0.5) / float(img_w)
v = (float(py) + 0.5) / float(img_h)
x = left + (right - left) * u
y = top + (bottom - top) * v
return wp.normalize(wp.vec3(x, y, -znear))
@wp.func
def pack_rgba_to_uint32(r: float, g: float, b: float, a: float) -> wp.uint32:
"""Pack RGBA values into a single uint32 for efficient memory access."""
return wp.uint32((int(a) << int(24)) | (int(r) << int(16)) | (int(g) << int(8)) | int(b))
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+48 -52
View File
@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src import smooth
from mujoco.mjx.third_party.mujoco_warp._src import support
from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import get_sdf_params
from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
@@ -42,7 +43,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@@ -124,10 +124,10 @@ def _magnetometer(
@wp.func
def _cam_projection(
# Model:
cam_fovy: wp.array(dtype=float),
cam_fovy: wp.array2d(dtype=float),
cam_resolution: wp.array(dtype=wp.vec2i),
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_intrinsic: wp.array(dtype=wp.vec4),
cam_intrinsic: wp.array2d(dtype=wp.vec4),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -138,8 +138,8 @@ def _cam_projection(
refid: int,
) -> wp.vec2:
sensorsize = cam_sensorsize[refid]
intrinsic = cam_intrinsic[refid]
fovy = cam_fovy[refid]
intrinsic = cam_intrinsic[worldid % cam_intrinsic.shape[0], refid]
fovy = cam_fovy[worldid % cam_fovy.shape[0], refid]
res = cam_resolution[refid]
target_xpos = site_xpos_in[worldid, objid]
@@ -470,10 +470,10 @@ def _sensor_pos(
site_quat: wp.array2d(dtype=wp.quat),
cam_bodyid: wp.array(dtype=int),
cam_quat: wp.array2d(dtype=wp.quat),
cam_fovy: wp.array(dtype=float),
cam_fovy: wp.array2d(dtype=float),
cam_resolution: wp.array(dtype=wp.vec2i),
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_intrinsic: wp.array(dtype=wp.vec4),
cam_intrinsic: wp.array2d(dtype=wp.vec4),
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objtype: wp.array(dtype=int),
@@ -782,7 +782,7 @@ def sensor_pos(m: Model, d: Data):
d,
rangefinder_pnt,
rangefinder_vec,
vec6(wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf),
vec6(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL),
True,
m.sensor_rangefinder_bodyid,
rangefinder_dist,
@@ -797,8 +797,7 @@ def sensor_pos(m: Model, d: Data):
energy_vel(m, d)
# collision sensors (distance, normal, fromto)
sensor_collision = wp.empty((d.nworld, m.nsensorcollision, 8, 7), dtype=float)
sensor_collision.fill_(1.0e32)
sensor_collision = wp.full((d.nworld, m.nsensorcollision, 8, 7), 1.0e32, dtype=float)
if m.nsensorcollision:
wp.launch(
_sensor_collision,
@@ -1787,38 +1786,12 @@ def _sensor_acc(
nmatch = sensor_contact_nmatch_in[worldid, contactsensorid]
if reduce == 3: # netforce
# compute point: force-weighted centroid of contact position
# Single-pass computation: first compute centroid, then wrench about centroid
# Pass 1: compute force-weighted centroid of contact positions
net_pos = wp.vec3(0.0)
total_force_magnitude = float(0.0)
for i in range(nmatch):
cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
contact_forcetorque = support.contact_force_fn(
opt_cone,
contact_frame_in,
contact_friction_in,
contact_dim_in,
contact_efc_address_in,
efc_force_in,
njmax_in,
nacon_in,
worldid,
cid,
False,
)
weight = wp.norm_l2(wp.spatial_top(contact_forcetorque))
net_pos += weight * contact_pos_in[cid]
total_force_magnitude += weight
net_pos /= wp.max(total_force_magnitude, MJ_MINVAL)
# TODO(team): iterate over matches once
# compute total wrench about point, in the global frame
net_force = wp.vec3(0.0)
net_torque = wp.vec3(0.0)
total_force_magnitude = float(0.0)
for i in range(nmatch):
cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
@@ -1837,8 +1810,15 @@ def _sensor_acc(
cid,
False,
)
contact_forcetorque *= dir
# Accumulate for centroid computation (unsigned force magnitude)
weight = wp.norm_l2(wp.spatial_top(contact_forcetorque))
contact_pos = contact_pos_in[cid]
net_pos += weight * contact_pos
total_force_magnitude += weight
# Apply direction and transform to global frame
contact_forcetorque *= dir
force_local = wp.spatial_top(contact_forcetorque)
torque_local = wp.spatial_bottom(contact_forcetorque)
@@ -1848,12 +1828,18 @@ def _sensor_acc(
force_global = frameT @ force_local
torque_global = frameT @ torque_local
# add to total force, torque
# Accumulate force and torque (about origin for now)
net_force += force_global
net_torque += torque_global
# Accumulate moment contribution: will adjust after centroid is computed
net_torque += wp.cross(contact_pos, force_global)
# add induced moment: torque += (pos - point) x force
net_torque += wp.cross(contact_pos_in[cid] - net_pos, force_global)
# Finalize centroid
net_pos /= wp.max(total_force_magnitude, MJ_MINVAL)
# Adjust torque: subtract moment from centroid (since we accumulated about origin)
# torque_about_centroid = torque_about_origin - centroid x total_force
net_torque -= wp.cross(net_pos, net_force)
adr_slot = adr
@@ -1888,7 +1874,8 @@ def _sensor_acc(
out[adr_slot + 1] = 1.0
out[adr_slot + 2] = 0.0
else:
for i in range(wp.min(nmatch, num)):
nslots = wp.min(nmatch, num)
for i in range(nslots):
# sorted contact id
cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
@@ -2165,7 +2152,15 @@ def _sensor_tactile(
contact_type = geom_type[geom]
plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params(
oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, contact_type, geom_size[worldid, geom], plugin_id, mesh_id
oct_child,
oct_aabb,
oct_coeff,
plugin,
plugin_attr,
contact_type,
geom_size[worldid % geom_size.shape[0], geom],
plugin_id,
mesh_id,
)
depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0)
@@ -2357,16 +2352,16 @@ def _contact_match(
@cache_kernel
def _contact_sort(maxmatch: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def contact_sort(
# Model:
sensor_intprm: wp.array2d(dtype=int),
sensor_contact_adr: wp.array(dtype=int),
# Data in:
# In:
sensor_contact_nmatch_in: wp.array2d(dtype=int),
sensor_contact_matchid_in: wp.array3d(dtype=int),
sensor_contact_criteria_in: wp.array3d(dtype=float),
# Data out:
# Out:
sensor_contact_matchid_out: wp.array3d(dtype=int),
):
worldid, contactsensorid = wp.tid()
@@ -2819,13 +2814,13 @@ def energy_pos(m: Model, d: Data):
@cache_kernel
def _energy_vel_kinetic(nv: int):
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def energy_vel_kinetic(
# Data in:
qvel_in: wp.array2d(dtype=float),
# In:
Mqvel: wp.array2d(dtype=float),
# Out:
# Data out:
energy_out: wp.array(dtype=wp.vec2),
):
worldid = wp.tid()
@@ -2849,12 +2844,13 @@ def energy_vel(m: Model, d: Data):
# kinetic energy: 0.5 * qvel.T @ M @ qvel
# M @ qvel
support.mul_m(m, d, d.efc.mv, d.qvel)
mv = wp.zeros((d.nworld, m.nv), dtype=float)
support.mul_m(m, d, mv, d.qvel)
wp.launch_tiled(
_energy_vel_kinetic(m.nv),
dim=d.nworld,
inputs=[d.qvel, d.efc.mv],
inputs=[d.qvel, mv],
outputs=[d.energy],
block_dim=m.block_dim.energy_vel_kinetic,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+384 -217
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+1432 -787
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+110 -85
View File
@@ -23,21 +23,21 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import State
from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
wp.set_module_options({"enable_backward": False})
@cache_kernel
def mul_m_sparse_diag(check_skip: bool):
@nested_kernel(module="unique", enable_backward=False)
def _mul_m_sparse_diag(
def mul_m_sparse(check_skip: bool):
@wp.kernel(module="unique")
def _mul_m_sparse(
# Model:
dof_Madr: wp.array(dtype=int),
qM_mulm_rowadr: wp.array(dtype=int),
qM_mulm_col: wp.array(dtype=int),
qM_mulm_madr: wp.array(dtype=int),
# Data in:
qM_in: wp.array3d(dtype=float),
# In:
@@ -46,26 +46,33 @@ def mul_m_sparse_diag(check_skip: bool):
# Out:
res: wp.array2d(dtype=float),
):
"""Diagonal update for sparse matmul."""
"""Sparse matmul: one thread per DOF, gather-based (no atomics)."""
worldid, dofid = wp.tid()
if wp.static(check_skip):
if skip[worldid]:
return
res[worldid, dofid] = qM_in[worldid, 0, dof_Madr[dofid]] * vec[worldid, dofid]
# Gather all contributions (diagonal + off-diagonal)
acc = float(0.0)
start = qM_mulm_rowadr[dofid]
end = qM_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]
return _mul_m_sparse_diag
res[worldid, dofid] = acc
return _mul_m_sparse
@cache_kernel
def mul_m_sparse_ij(check_skip: bool):
@nested_kernel(module="unique", enable_backward=False)
def _mul_m_sparse_ij(
# Model:
qM_mulm_i: wp.array(dtype=int),
qM_mulm_j: wp.array(dtype=int),
qM_madr_ij: wp.array(dtype=int),
def mul_m_dense(nv: int, check_skip: bool):
"""Simple SIMT dense matmul: one thread per output element."""
@wp.kernel(module="unique")
def _mul_m_dense(
# Data in:
qM_in: wp.array3d(dtype=float),
# In:
@@ -74,52 +81,16 @@ def mul_m_sparse_ij(check_skip: bool):
# Out:
res: wp.array2d(dtype=float),
):
"""Off-diagonal update for sparse matmul."""
worldid, elementid = wp.tid()
worldid, i = wp.tid()
if wp.static(check_skip):
if skip[worldid]:
return
i = qM_mulm_i[elementid]
j = qM_mulm_j[elementid]
madr_ij = qM_madr_ij[elementid]
qM_ij = qM_in[worldid, 0, madr_ij]
wp.atomic_add(res[worldid], i, qM_ij * vec[worldid, j])
wp.atomic_add(res[worldid], j, qM_ij * vec[worldid, i])
return _mul_m_sparse_ij
@cache_kernel
def mul_m_dense(tile: TileSet, check_skip: bool):
"""Returns a matmul kernel for some tile size."""
@nested_kernel(module="unique", enable_backward=False)
def _mul_m_dense(
# Data In:
qM_in: wp.array3d(dtype=float),
# In:
adr: wp.array(dtype=int),
vec: wp.array3d(dtype=float),
skip: wp.array(dtype=bool),
# Out:
res: wp.array3d(dtype=float),
):
worldid, nodeid = wp.tid()
TILE_SIZE = wp.static(tile.size)
if wp.static(check_skip):
if skip[worldid]:
return
dofid = adr[nodeid]
qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False)
vec_tile = wp.tile_load(vec[worldid], shape=(TILE_SIZE, 1), offset=(dofid, 0), bounds_check=False)
res_tile = wp.tile_matmul(qM_tile, vec_tile)
wp.tile_store(res[worldid], res_tile, offset=(dofid, 0), bounds_check=False)
acc = float(0.0)
for j in range(wp.static(nv)):
acc += qM_in[worldid, i, j] * vec[worldid, j]
res[worldid, i] = acc
return _mul_m_dense
@@ -149,36 +120,21 @@ def mul_m(
if M is None:
M = d.qM
if m.opt.is_sparse:
if m.is_sparse:
wp.launch(
mul_m_sparse_diag(check_skip),
mul_m_sparse(check_skip),
dim=(d.nworld, m.nv),
inputs=[m.dof_Madr, M, vec, skip],
outputs=[res],
)
wp.launch(
mul_m_sparse_ij(check_skip),
dim=(d.nworld, m.qM_madr_ij.size),
inputs=[m.qM_mulm_i, m.qM_mulm_j, m.qM_madr_ij, M, vec, skip],
inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
outputs=[res],
)
else:
for tile in m.qM_tiles:
wp.launch_tiled(
mul_m_dense(tile, check_skip),
dim=(d.nworld, tile.adr.size),
inputs=[
M,
tile.adr,
# note reshape: tile_matmul expects 2d input
vec.reshape(vec.shape + (1,)),
skip,
],
outputs=[res.reshape(res.shape + (1,))],
block_dim=m.block_dim.mul_m_dense,
)
wp.launch(
mul_m_dense(m.nv, check_skip),
dim=(d.nworld, m.nv),
inputs=[M, vec, skip],
outputs=[res],
)
@wp.kernel
@@ -408,7 +364,7 @@ def transform_force(frc: wp.spatial_vector, offset: wp.vec3) -> wp.spatial_vecto
@wp.func
def jac(
def jac_dof(
# Model:
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
@@ -446,8 +402,77 @@ def jac(
return jacp, jacr
@cache_kernel
def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
@wp.kernel(module="unique", enable_backward=False)
def _jac(
# Model:
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
dof_bodyid: wp.array(dtype=int),
# Data in:
subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector),
# In:
point_in: wp.array(dtype=wp.vec3),
bodyid_in: wp.array(dtype=int),
# Out:
jacp_out: wp.array3d(dtype=float),
jacr_out: wp.array3d(dtype=float),
):
worldid, dofid = wp.tid()
jacp_val, jacr_val = jac_dof(
body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, point_in[worldid], bodyid_in[worldid], dofid, worldid
)
if wp.static(has_jacp):
jacp_out[worldid, 0, dofid] = jacp_val[0]
jacp_out[worldid, 1, dofid] = jacp_val[1]
jacp_out[worldid, 2, dofid] = jacp_val[2]
if wp.static(has_jacr):
jacr_out[worldid, 0, dofid] = jacr_val[0]
jacr_out[worldid, 1, dofid] = jacr_val[1]
jacr_out[worldid, 2, dofid] = jacr_val[2]
return _jac
@event_scope
def jac(
m: Model,
d: Data,
jacp: wp.array | None, # wp.array3d(dtype=float)
jacr: wp.array | None, # wp.array3d(dtype=float)
point: wp.array(dtype=wp.vec3),
body: wp.array(dtype=int),
):
"""Compute translational and rotational Jacobian for point on body.
Args:
m: The model containing kinematic and dynamic information (device).
d: The data object containing the current state (device).
jacp: Output translational Jacobian (optional).
jacr: Output rotational Jacobian (optional).
point: 3D point in global coordinates.
body: Body ID for each world.
"""
kernel = _make_jac_kernel(jacp is not None, jacr is not None)
jacp_arr = jacp or wp.empty((0, 0, 0), dtype=float)
jacr_arr = jacr or wp.empty((0, 0, 0), dtype=float)
wp.launch(
kernel,
dim=(d.nworld, m.nv),
inputs=[m.body_parentid, m.body_rootid, m.dof_bodyid, d.subtree_com, d.cdof, point, body],
outputs=[jacp_arr, jacr_arr],
)
@wp.func
def jac_dot(
def jac_dot_dof(
# Model:
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
@@ -529,7 +554,7 @@ def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
if sig >= (1 << State.NSTATE):
raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def _get_state(
# Model:
nq: int,
@@ -668,7 +693,7 @@ def set_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
if sig >= (1 << State.NSTATE):
raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
@nested_kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False)
def _set_state(
# Model:
nq: int,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+216 -85
View File
@@ -32,6 +32,9 @@ MJ_MAX_EPAFACES = 5
TILE_SIZE_JTDAJ_SPARSE = 16
TILE_SIZE_JTDAJ_DENSE = 16
# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
TEXTURE_DTYPE = wp.Texture2D if hasattr(wp, "Texture2D") else int
# TODO(team): add check that all wp.launch_tiled 'block_dim' settings are configurable
@dataclasses.dataclass
@@ -61,9 +64,9 @@ class BlockDim:
update_gradient_cholesky_blocked: int = 32
update_gradient_JTDAJ_sparse: int = 64
update_gradient_JTDAJ_dense: int = 96
linesearch_iterative: int = 64
# support
mul_m_dense: int = 32
linesearch_iterative: int = 32
# derivative
qderiv_actuator_dense: int = 32
class BroadphaseType(enum.IntEnum):
@@ -114,6 +117,23 @@ class CamLightType(enum.IntEnum):
TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM
class ProjectionType(enum.IntEnum):
"""Type of camera projection.
Attributes:
PERSPECTIVE: perspective projection
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
class DataType(enum.IntFlag):
"""Sensor data types.
@@ -164,7 +184,8 @@ class DisableBit(enum.IntFlag):
REFSAFE = mujoco.mjtDisableBit.mjDSBL_REFSAFE
SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR
EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP
# unsupported: MIDPHASE, AUTORESET, NATIVECCD, ISLAND
NATIVECCD = mujoco.mjtDisableBit.mjDSBL_NATIVECCD
# unsupported: MIDPHASE, AUTORESET, ISLAND
class EnableBit(enum.IntFlag):
@@ -321,6 +342,20 @@ class GeomType(enum.IntEnum):
# unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
class CollisionType(enum.IntEnum):
"""Type of narrowphase collision.
Attributes:
PRIMITIVE: primitive collision
CONVEX: convex collision (CCD)
SDF: sdf collision
"""
PRIMITIVE = 0
CONVEX = 1
SDF = 2
class SolverType(enum.IntEnum):
"""Constraint solver algorithm.
@@ -665,10 +700,8 @@ class Option:
warp only fields:
impratio_invsqrt: ratio of friction-to-normal contact impedance (stored as inverse square root)
is_sparse: whether to use sparse representations
ls_parallel: evaluate engine solver step sizes in parallel
ls_parallel_min_step: minimum step size for solver linesearch
has_fluid: True if wind, density, or viscosity are non-zero at put_model time
broadphase: broadphase type (BroadphaseType)
broadphase_filter: broadphase filter bitflag (BroadphaseFilter)
graph_conditional: flag to use cuda graph conditional
@@ -700,10 +733,8 @@ class Option:
sdf_iterations: int
# warp only fields:
impratio_invsqrt: array("*", float)
is_sparse: bool
ls_parallel: bool
ls_parallel_min_step: float
has_fluid: bool
broadphase: BroadphaseType
broadphase_filter: BroadphaseFilter
graph_conditional: bool
@@ -716,10 +747,10 @@ class Statistic:
"""Model statistics (in qpos0).
Attributes:
meaninertia: mean diagonal inertia
meaninertia: mean diagonal inertia (per-world)
"""
meaninertia: float
meaninertia: array("*", float)
@dataclasses.dataclass
@@ -749,6 +780,7 @@ class Model:
nbody: number of bodies
noct: number of total octree cells in all meshes
njnt: number of joints
ntree: number of kinematic trees
nM: number of non-zeros in sparse inertia matrix
nC: number of non-zeros in sparse body-dof matrix
ngeom: number of geoms
@@ -794,6 +826,7 @@ class Model:
body_jntadr: start addr of joints; -1: no joints (nbody,)
body_dofnum: number of motion degrees of freedom (nbody,)
body_dofadr: start addr of dofs; -1: no dofs (nbody,)
body_treeid: id of body's tree; -1: static (nbody,)
body_geomnum: number of geoms (nbody,)
body_geomadr: start addr of geoms; -1: no geoms (nbody,)
body_pos: position offset rel. to parent body (*, nbody, 3)
@@ -828,6 +861,7 @@ class Model:
dof_bodyid: id of dof's body (nv,)
dof_jntid: id of dof's joint (nv,)
dof_parentid: id of dof's parent; -1: none (nv,)
dof_treeid: id of dof's tree (nv,)
dof_Madr: dof address in M-diagonal (nv,)
dof_solref: constraint solver reference: frictionloss (*, nv, NREF)
dof_solimp: constraint solver impedance: frictionloss (*, nv, NIMP)
@@ -835,6 +869,9 @@ class Model:
dof_armature: dof armature inertia/mass (*, nv)
dof_damping: damping coefficient (*, nv)
dof_invweight0: diag. inverse inertia in qpos0 (*, nv)
tree_bodynum: number of bodies in tree (incl. root) (ntree,)
tree_dofadr: start address of tree's dofs (ntree,)
tree_dofnum: number of dofs in tree (ntree,)
geom_type: geometric type (GeomType) (ngeom,)
geom_contype: geom contact type (ngeom,)
geom_conaffinity: geom contact affinity (ngeom,)
@@ -870,10 +907,11 @@ class Model:
cam_poscom0: global position rel. to sub-com in qpos0 (*, ncam, 3)
cam_pos0: global position rel. to body in qpos0 (*, ncam, 3)
cam_mat0: global orientation in qpos0 (*, ncam, 3, 3)
cam_fovy: y field-of-view (ortho ? len : deg) (ncam,)
cam_projection: projection type (ProjectionType) (ncam,)
cam_fovy: y field-of-view (ortho ? len : deg) (*, ncam)
cam_resolution: resolution: pixels [width, height] (ncam, 2)
cam_sensorsize: sensor size: length [width, height] (ncam, 2)
cam_intrinsic: [focal length; principal point] (ncam, 4)
cam_intrinsic: [focal length; principal point] (*, ncam, 4)
light_mode: light tracking mode (CamLightType) (nlight,)
light_bodyid: id of light's body (nlight,)
light_targetbodyid: id of targeted body; -1: none (nlight,)
@@ -903,6 +941,9 @@ class Model:
flex_stiffness: finite element stiffness matrix (nflexelem, 21)
flex_bending: bending stiffness (nflexedge, 17)
flex_damping: Rayleigh's damping coefficient (nflex,)
flexedge_J_rownnz: number of nonzeros in Jacobian row (nflexedge,)
flexedge_J_rowadr: row start address in colind array (nflexedge,)
flexedge_J_colind: column indices in sparse Jacobian (nJfe,)
mesh_vertadr: first vertex address (nmesh,)
mesh_vertnum: number of vertices (nmesh,)
mesh_faceadr: first face address (nmesh,)
@@ -927,7 +968,7 @@ class Model:
hfield_ncol: number of columns in grid (nhfield,)
hfield_adr: start address in hfield_data (nhfield,)
hfield_data: elevation data (nhfielddata,)
mat_texid: texture id for rendering (nmat, mjNTEXROLE)
mat_texid: texture id for rendering (*, nmat, mjNTEXROLE)
mat_texrepeat: texture repeat for rendering (*, nmat, 2)
mat_rgba: rgba (*, nmat, 4)
pair_dim: contact dimensionality (npair,)
@@ -1008,6 +1049,7 @@ class Model:
mapM2M: index mapping from M (legacy) to M (CSR) (nC)
warp only fields:
nbranch: number of branches (leaf-to-root paths)
nv_pad: number of degrees of freedom + padding
nacttrnbody: number of actuators with body transmission
nsensorcollision: number of unique collisions for
@@ -1019,9 +1061,13 @@ class Model:
nmaxpyramid: maximum number of pyramid directions
nmaxpolygon: maximum number of verts per polygon
nmaxmeshdeg: maximum number of polygons per vert
is_sparse: whether to use sparse representations
has_fluid: True if wind, density, or viscosity are non-zero at put_model time
has_sdf_geom: whether the model contains SDF geoms
block_dim: block dim options
body_tree: list of body ids by tree level
body_branches: flattened body ids for all branches
body_branch_start: start index in body_branches for each branch (nbranch + 1,)
mocap_bodyid: id of body for mocap (nmocap,)
body_fluid_ellipsoid: does body use ellipsoid fluid (nbody,)
jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
@@ -1085,9 +1131,9 @@ class Model:
qLD_updates: tuple of index triples for sparse factorization
qM_fullm_i: sparse mass matrix addressing
qM_fullm_j: sparse mass matrix addressing
qM_mulm_i: sparse matmul addressing
qM_mulm_j: sparse matmul addressing
qM_madr_ij: sparse matmul addressing
qM_mulm_rowadr: sparse matmul row pointers
qM_mulm_col: sparse matmul column indices
qM_mulm_madr: sparse matmul matrix addresses
"""
nq: int
@@ -1097,6 +1143,7 @@ class Model:
nbody: int
noct: int
njnt: int
ntree: int
nM: int
nC: int
ngeom: int
@@ -1142,6 +1189,7 @@ class Model:
body_jntadr: array("nbody", int)
body_dofnum: array("nbody", int)
body_dofadr: array("nbody", int)
body_treeid: array("nbody", int)
body_geomnum: array("nbody", int)
body_geomadr: array("nbody", int)
body_pos: array("*", "nbody", wp.vec3)
@@ -1176,6 +1224,7 @@ class Model:
dof_bodyid: array("nv", int)
dof_jntid: array("nv", int)
dof_parentid: array("nv", int)
dof_treeid: array("nv", int)
dof_Madr: array("nv", int)
dof_solref: array("*", "nv", wp.vec2)
dof_solimp: array("*", "nv", vec5)
@@ -1183,6 +1232,9 @@ class Model:
dof_armature: array("*", "nv", float)
dof_damping: array("*", "nv", float)
dof_invweight0: array("*", "nv", float)
tree_bodynum: array("ntree", int)
tree_dofadr: array("ntree", int)
tree_dofnum: array("ntree", int)
geom_type: array("ngeom", int)
geom_contype: array("ngeom", int)
geom_conaffinity: array("ngeom", int)
@@ -1218,10 +1270,11 @@ class Model:
cam_poscom0: array("*", "ncam", wp.vec3)
cam_pos0: array("*", "ncam", wp.vec3)
cam_mat0: array("*", "ncam", wp.mat33)
cam_fovy: array("ncam", float)
cam_projection: array("ncam", int)
cam_fovy: array("*", "ncam", float)
cam_resolution: array("ncam", wp.vec2i)
cam_sensorsize: array("ncam", wp.vec2)
cam_intrinsic: array("ncam", wp.vec4)
cam_intrinsic: array("*", "ncam", wp.vec4)
light_mode: array("nlight", int)
light_bodyid: array("nlight", int)
light_targetbodyid: array("nlight", int)
@@ -1251,6 +1304,9 @@ class Model:
flex_stiffness: array("nflexelem", 21, float)
flex_bending: array("nflexedge", 17, float)
flex_damping: array("nflex", float)
flexedge_J_rownnz: array("nflexedge", int)
flexedge_J_rowadr: array("nflexedge", int)
flexedge_J_colind: wp.array(dtype=int)
mesh_vertadr: array("nmesh", int)
mesh_vertnum: array("nmesh", int)
mesh_faceadr: array("nmesh", int)
@@ -1275,7 +1331,7 @@ class Model:
hfield_ncol: array("nhfield", int)
hfield_adr: array("nhfield", int)
hfield_data: array("nhfielddata", float)
mat_texid: array("nmat", 10, int)
mat_texid: array("*", "nmat", 10, int)
mat_texrepeat: array("*", "nmat", wp.vec2)
mat_rgba: array("*", "nmat", wp.vec4)
pair_dim: array("npair", int)
@@ -1355,6 +1411,7 @@ class Model:
M_colind: array("nC", int)
mapM2M: array("nC", int)
# warp only fields:
nbranch: int
nv_pad: int
nacttrnbody: int
nsensorcollision: int
@@ -1365,9 +1422,13 @@ class Model:
nmaxpyramid: int
nmaxpolygon: int
nmaxmeshdeg: int
is_sparse: bool
has_fluid: bool
has_sdf_geom: bool
block_dim: BlockDim
body_tree: tuple[wp.array(dtype=int), ...]
body_branches: wp.array(dtype=int)
body_branch_start: wp.array(dtype=int)
mocap_bodyid: array("nmocap", int)
body_fluid_ellipsoid: array("nbody", bool)
jnt_limited_slide_hinge_adr: wp.array(dtype=int)
@@ -1422,9 +1483,10 @@ class Model:
qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...]
qM_fullm_i: wp.array(dtype=int)
qM_fullm_j: wp.array(dtype=int)
qM_mulm_i: wp.array(dtype=int)
qM_mulm_j: wp.array(dtype=int)
qM_madr_ij: wp.array(dtype=int)
# Gather-based sparse mul_m indices (thread per DOF, no atomics)
qM_mulm_rowadr: wp.array(dtype=int) # start address for each row [nv+1]
qM_mulm_col: wp.array(dtype=int) # column index to gather from
qM_mulm_madr: wp.array(dtype=int) # matrix address to read
class ContactType(enum.IntFlag):
@@ -1492,26 +1554,9 @@ class Constraint:
aref: reference pseudo-acceleration (nworld, njmax)
frictionloss: frictionloss (friction) (nworld, njmax)
force: constraint force in constraint space (nworld, njmax)
Jaref: Jac*qacc - aref (nworld, njmax)
Ma: M*qacc (nworld, nv)
grad: gradient of master cost (nworld, nv_pad)
grad_dot: dot(grad, grad) (nworld,)
Mgrad: M / grad (nworld, nv_pad)
search: linesearch vector (nworld, nv)
search_dot: dot(search, search) (nworld,)
gauss: Gauss Cost (nworld,)
cost: constraint + Gauss cost (nworld,)
prev_cost: cost from previous iter (nworld,)
state: constraint state (nworld, njmax_pad)
mv: qM @ search (nworld, nv)
jv: efc_J @ search (nworld, njmax)
quad: quadratic cost coefficients (nworld, njmax, 3)
quad_gauss: quadratic cost Gauss coefficients (nworld, 3)
alpha: line search step size (nworld,)
prev_grad: previous grad (nworld, nv)
prev_Mgrad: previous Mgrad (nworld, nv)
beta: Polak-Ribiere beta (nworld,)
done: solver done (nworld,)
warp only fields:
Ma: M*qacc (nworld, nv)
"""
type: array("nworld", "njmax", int)
@@ -1524,26 +1569,8 @@ class Constraint:
aref: array("nworld", "njmax", float)
frictionloss: array("nworld", "njmax", float)
force: array("nworld", "njmax", float)
Jaref: array("nworld", "njmax", float)
Ma: array("nworld", "nv", float)
grad: array("nworld", "nv_pad", float)
grad_dot: array("nworld", float)
Mgrad: array("nworld", "nv_pad", float)
search: array("nworld", "nv", float)
search_dot: array("nworld", float)
gauss: array("nworld", float)
cost: array("nworld", float)
prev_cost: array("nworld", float)
state: array("nworld", "njmax_pad", int)
mv: array("nworld", "nv", float)
jv: array("nworld", "njmax", float)
quad: array("nworld", "njmax", wp.vec3)
quad_gauss: array("nworld", wp.vec3)
alpha: array("nworld", float)
prev_grad: array("nworld", "nv", float)
prev_Mgrad: array("nworld", "nv", float)
beta: array("nworld", float)
done: array("nworld", bool)
Ma: array("nworld", "nv", float)
@dataclasses.dataclass
@@ -1590,7 +1617,7 @@ class Data:
cdof: com-based motion axis of each dof (rot:lin) (nworld, nv, 6)
cinert: com-based body inertia and mass (nworld, nbody, 10)
flexvert_xpos: cartesian flex vertex positions (nworld, nflexvert, 3)
flexedge_J: edge length Jacobian (nworld, nflexedge, nv)
flexedge_J: edge length Jacobian (nworld, 1, nflexedge*6)
flexedge_length: flex edge lengths (nworld, nflexedge, 1)
ten_wrapadr: start address of tendon's path (nworld, ntendon)
ten_wrapnum: number of wrap points in path (nworld, ntendon)
@@ -1606,7 +1633,7 @@ class Data:
qLD: L'*D*L factorization of M (nworld, nv, nv) if dense
(nworld, 1, nC) if sparse
qLDiagInv: 1/diag(D) (nworld, nv)
flexedge_velocity: flex edge velocities (nworld, nflexedge,)
flexedge_velocity: flex edge velocities (nworld, nflexedge)
ten_velocity: tendon velocities (nworld, ntendon)
actuator_velocity: actuator velocities (nworld, nu)
cvel: com-based velocity (rot:lin) (nworld, nbody, 6)
@@ -1636,18 +1663,9 @@ class Data:
warp only fields:
nworld: number of worlds
naconmax: maximum number of contacts (shared across all worlds)
naccdmax: maximum number of contacts for CCD (all worlds)
njmax: maximum number of constraints per world
nacon: number of detected contacts (across all worlds) (1,)
ne_connect: number of equality connect constraints (nworld,)
ne_weld: number of equality weld constraints (nworld,)
ne_jnt: number of equality joint constraints (nworld,)
ne_ten: number of equality tendon constraints (nworld,)
ne_flex: number of flex edge equality constraints (nworld,)
nsolving: number of unconverged worlds (1,)
subtree_bodyvel: subtree body velocity (ang, vel) (nworld, nbody, 6)
collision_pair: collision pairs from broadphase (naconmax, 2)
collision_pairid: ids from broadphase (naconmax, 2)
collision_worldid: collision world ids from broadphase (naconmax,)
ncollision: collision count from broadphase (1,)
"""
@@ -1690,7 +1708,7 @@ class Data:
cdof: array("nworld", "nv", wp.spatial_vector)
cinert: array("nworld", "nbody", vec10)
flexvert_xpos: array("nworld", "nflexvert", wp.vec3)
flexedge_J: array("nworld", "nflexedge", "nv", float)
flexedge_J: wp.array3d(dtype=float)
flexedge_length: array("nworld", "nflexedge", float)
ten_wrapadr: array("nworld", "ntendon", int)
ten_wrapnum: array("nworld", "ntendon", int)
@@ -1732,18 +1750,131 @@ class Data:
# warp only fields:
nworld: int
naconmax: int
naccdmax: int
njmax: int
nacon: array(1, int)
ne_connect: array("nworld", int)
ne_weld: array("nworld", int)
ne_jnt: array("nworld", int)
ne_ten: array("nworld", int)
ne_flex: array("nworld", int)
nsolving: array(1, int)
subtree_bodyvel: array("nworld", "nbody", wp.spatial_vector)
# warp only: collision driver
collision_pair: array("naconmax", wp.vec2i)
collision_pairid: array("naconmax", wp.vec2i)
collision_worldid: array("naconmax", int)
ncollision: array(1, int)
@dataclasses.dataclass
class CollisionContext:
"""Collision driver intermediate arrays.
Attributes:
collision_pair: collision pairs from broadphase (naconmax, 2)
collision_pairid: ids from broadphase (naconmax, 2)
collision_worldid: collision world ids from broadphase (naconmax,)
"""
collision_pair: wp.array
collision_pairid: wp.array
collision_worldid: wp.array
@dataclasses.dataclass
class RenderContext:
"""Context for rendering.
Attributes:
nrender: number of actively rendering cameras
cam_res: camera resolution for actively rendering cameras
cam_id_map: camera id map
use_textures: whether to use textures
use_shadows: whether to use shadows
bvh_ngeom: number of geometries in the BVH
enabled_geom_ids: enabled geometry ids
mesh_registry: mesh BVH id to warp mesh mapping
mesh_bvh_id: mesh BVH ids
mesh_bounds_size: mesh bounds size
mesh_texcoord: mesh texture coordinates
mesh_texcoord_offsets: mesh texture coordinate offsets
mesh_facetexcoord: mesh face texture coordinates
textures: textures
textures_registry: texture registry
hfield_registry: hfield BVH id to warp mesh mapping
hfield_bvh_id: hfield BVH ids
hfield_bounds_size: hfield bounds size
flex_mesh: flex mesh
flex_rgba: flex rgba
flex_bvh_id: flex BVH id
flex_face_point: flex face points
flex_faceadr: flex face addresses
flex_nface: number of flex faces
flex_nwork: total flex work items for refit
flex_group_root: flex group roots
flex_elemdataadr: flex element data addresses
flex_shell: flex shell data
flex_shelldataadr: flex shell data addresses
flex_radius: flex radius
flex_workadr: flex work item addresses for refit
flex_worknum: flex work item counts for refit
flex_render_smooth: whether to render flex meshes smoothly
bvh: scene BVH
bvh_id: scene BVH id
lower: lower bounds
upper: upper bounds
group: groups
group_root: group roots
ray: rays
rgb_data: RGB data
rgb_adr: RGB addresses
rgb_size: per-camera RGB buffer sizes
depth_data: depth data
depth_adr: depth addresses
depth_size: per-camera depth buffer sizes
render_rgb: per-camera RGB render flags
render_depth: per-camera depth render flags
znear: near plane distance
total_rays: total number of rays
"""
nrender: int
cam_res: array("ncam", wp.vec2i)
cam_id_map: array("ncam", int)
use_textures: bool
use_shadows: bool
background_color: wp.uint32
bvh_ngeom: int
enabled_geom_ids: array("*", int)
mesh_registry: dict
mesh_bvh_id: array("nmesh", wp.uint64)
mesh_bounds_size: array("nmesh", wp.vec3)
mesh_texcoord: array("*", wp.vec2)
mesh_texcoord_offsets: array("nmesh", int)
mesh_facetexcoord: array("nmeshface", wp.vec3i)
# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
textures: array("*", TEXTURE_DTYPE)
textures_registry: list[TEXTURE_DTYPE]
hfield_registry: dict
hfield_bvh_id: array("nhfield", wp.uint64)
hfield_bounds_size: array("nhfield", wp.vec3)
flex_mesh: wp.Mesh
flex_rgba: array("nflex", wp.vec4)
flex_bvh_id: wp.uint64
flex_face_point: array("*", wp.vec3)
flex_faceadr: array("nflex", int)
flex_nface: int
flex_nwork: int
flex_group_root: array("nworld", int)
flex_elemdataadr: array("nflex", int)
flex_shell: array("*", int)
flex_shelldataadr: array("nflex", int)
flex_radius: array("nflex", float)
flex_workadr: array("nflex", int)
flex_worknum: array("nflex", int)
flex_render_smooth: bool
bvh: wp.Bvh
bvh_id: wp.uint64
lower: array("*", wp.vec3)
upper: array("*", wp.vec3)
group: array("*", int)
group_root: array("*", int)
ray: array("*", wp.vec3)
rgb_data: array("*", wp.uint32)
rgb_adr: array("ncam", int)
depth_data: array("*", wp.float32)
depth_adr: array("ncam", int)
render_rgb: array("ncam", bool)
render_depth: array("ncam", bool)
znear: float
total_rays: int
mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_misc.py
Vendored
+15 -14
View File
@@ -20,6 +20,7 @@ from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import math
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 GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
@@ -105,13 +106,13 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
Args:
end: Two 2D points.
side: Optional 2D side point, no side point: wp.vec2(wp.inf).
side: Optional 2D side point, no side point: wp.vec2(MJ_MAXVAL).
radius: Circle radius.
Returns:
Length of circular wrap or -1.0 if no wrap, pair of 2D wrap points.
"""
valid_side = wp.norm_l2(side) < wp.inf
valid_side = wp.norm_l2(side) < MJ_MAXVAL
end0 = wp.vec2(end[0], end[1])
end1 = wp.vec2(end[2], end[3])
@@ -122,13 +123,13 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
# either point inside circle or circle too small: no wrap
if (sqlen0 < sqrad) or (sqlen1 < sqrad) or (radius < MJ_MINVAL):
return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
# points too close: no wrap
dif = end1 - end0
dd = wp.dot(dif, dif)
if dd < MJ_MINVAL:
return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
# find nearest point on line segment to origin: a * dif + d0
a = -wp.dot(dif, end0) / dd
@@ -137,7 +138,7 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
# check for intersection and side
tmp = a * dif + end0
if (wp.dot(tmp, tmp) > sqrad) and (not valid_side or wp.dot(side, tmp) >= 0.0):
return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
sqrt0 = wp.sqrt(sqlen0 - sqrad)
sqrt1 = wp.sqrt(sqlen1 - sqrad)
@@ -191,7 +192,7 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
# check for intersection
if is_intersect(end0, pnt0, end1, pnt1):
return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
# return curve length
return length_circle(pnt0, pnt1, ind, radius), pnt0, pnt1
@@ -230,7 +231,7 @@ def wrap_inside(
# either point inside circle or circle too small: no wrap
if (len0 <= radius) or (len1 <= radius) or (radius < MJ_MINVAL) or (len0 < MJ_MINVAL) or (len1 < MJ_MINVAL):
return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
# segment-circle intersection: no wrap
if dd > MJ_MINVAL:
@@ -241,7 +242,7 @@ def wrap_inside(
if (a > 0.0) and (a < 1.0):
tmp = end0 + a * dif
if wp.norm_l2(tmp) <= radius:
return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
# prepare default in case of numerical failure: average
pnt = 0.5 * (end0 + end1)
@@ -335,14 +336,14 @@ def wrap(
mat: Orientation of geom.
radius: Geom radius.
geomtype: Wrap type (mjtWrap).
side: 3D position for sidesite, no side point: wp.vec3(wp.inf).
side: 3D position for sidesite, no side point: wp.vec3(MJ_MAXVAL).
Returns:
Length of circular wrap else -1.0 if no wrap, pair of 3D wrap points.
"""
# check object type
if geomtype != WrapType.SPHERE and geomtype != WrapType.CYLINDER:
return wp.inf, wp.vec3(wp.inf), wp.vec3(wp.inf)
return MJ_MAXVAL, wp.vec3(MJ_MAXVAL), wp.vec3(MJ_MAXVAL)
# map sites to wrap object's local frame
matT = wp.transpose(mat)
@@ -351,7 +352,7 @@ def wrap(
# too close to origin: return
if (wp.norm_l2(p0) < MJ_MINVAL) or (wp.norm_l2(p1) < MJ_MINVAL):
return -1.0, wp.vec3(wp.inf), wp.vec3(wp.inf)
return -1.0, wp.vec3(MJ_MAXVAL), wp.vec3(MJ_MAXVAL)
# construct 2D frame for circle wrap
if geomtype == WrapType.SPHERE:
@@ -399,7 +400,7 @@ def wrap(
)
# handle sidesite
valid_side = wp.norm_l2(side) < wp.inf
valid_side = wp.norm_l2(side) < MJ_MAXVAL
if valid_side:
# side point: apply same projection as x0, x1
@@ -414,7 +415,7 @@ def wrap(
sidepnt_proj, _ = math.normalize_with_norm(sidepnt_proj)
sidepnt_proj *= radius
else:
sidepnt_proj = wp.vec2(wp.inf)
sidepnt_proj = wp.vec2(MJ_MAXVAL)
# apply inside wrap
if valid_side and wp.norm_l2(sidepnt) < radius:
@@ -424,7 +425,7 @@ def wrap(
# no wrap: return
if wlen < 0.0:
return -1.0, wp.vec3(wp.inf), wp.vec3(wp.inf)
return -1.0, wp.vec3(MJ_MAXVAL), wp.vec3(MJ_MAXVAL)
# reconstruct 3D points in local frame: res
res0 = axis0 * pnt0[0] + axis1 * pnt0[1]
mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
Vendored
+1 -78
View File
@@ -15,7 +15,6 @@
import functools
import inspect
from typing import Callable, Optional
import warp as wp
@@ -119,82 +118,6 @@ def event_scope(fn, name: str = ""):
return wrapper
# @nested_kernel decorator to automatically set up modules based on nested
# function names
def nested_kernel(
f: Optional[Callable] = None,
*,
enable_backward: Optional[bool] = None,
module: Optional[wp.Module] = None,
):
"""Decorator to register a Warp kernel from a Python function.
The function must be defined with type annotations for all arguments.
The function must not return anything.
Example::
@nested_kernel
def my_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
tid = wp.tid()
b[tid] = a[tid] + 1.0
@nested_kernel(enable_backward=False)
def my_kernel_no_backward(a: wp.array(dtype=float, ndim=2), x: float):
# the backward pass will not be generated
i, j = wp.tid()
a[i, j] = x
@nested_kernel(module="unique")
def my_kernel_unique_module(a: wp.array(dtype=float), b: wp.array(dtype=float)):
# the kernel will be registered in new unique module created just for this
# kernel and its dependent functions and structs
tid = wp.tid()
b[tid] = a[tid] + 1.0
@neste_kernel(enable_backward=False, module=None)
def my_kernel_with_args(a: wp.array(dtype=float), b: wp.array(dtype=float)):
# can now use arguments even when module=None
tid = wp.tid()
b[tid] = a[tid] + 1.0
Args:
f: The function to be registered as a kernel.
enable_backward: If False, the backward pass will not be generated.
module: The :class:`warp.Module` to which the kernel belongs. Alternatively,
if a string `"unique"` is provided, the kernel is assigned to a new module
named after the kernel name and hash. If None, the module is inferred from
the function's module.
Returns:
The registered kernel.
"""
def decorator(func):
if module is None:
# create a module name based on the name of the nested function
# get the qualified name, e.g. "main.<locals>.nested_kernel"
qualname = func.__qualname__
parts = [part for part in qualname.split(".") if part != "<locals>"]
outer_functions = parts[:-1]
module_name = wp.get_module(".".join([func.__module__] + outer_functions))
else:
module_name = module
return wp.kernel(func, enable_backward=enable_backward, module=module_name)
# Handle both @kernel and @kernel(...) usage patterns
if f is None:
# Called with arguments: @kernel(enable_backward=False)
return decorator
else:
# Called without arguments: @kernel
return decorator(f)
_KERNEL_CACHE = {}
@@ -221,4 +144,4 @@ def check_toolkit_driver():
wp.init()
if wp.get_device().is_cuda:
if not wp.is_conditional_graph_supported():
RuntimeError("Minimum supported CUDA version: 12.4.")
raise RuntimeError("Minimum supported CUDA version: 12.4.")
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+1 -1
View File
@@ -54,7 +54,7 @@ dev = [
"ruff",
"pygls>=1.0.0,<2.0.0",
"lsprotocol>=2023.0.1,<2024.0.0",
"mujoco>=3.3.7.dev0",
"mujoco>=3.4.1.dev0",
"warp-lang>=1.11.0.dev0",
]
# TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+39 -30
View File
@@ -24,6 +24,7 @@ Example:
import copy
import enum
import logging
import shutil
import sys
import time
from typing import Sequence
@@ -51,10 +52,11 @@ class EngineOptions(enum.IntEnum):
C = 1
_CLEAR_KERNEL_CACHE = flags.DEFINE_bool("clear_kernel_cache", False, "Clear kernel cache (to calculate full JIT time)")
_CLEAR_WARP_CACHE = flags.DEFINE_bool("clear_warp_cache", False, "Clear warp caches (kernel, LTO, CUDA compute)")
_ENGINE = flags.DEFINE_enum_class("engine", EngineOptions.WARP, EngineOptions, "Simulation engine")
_NCONMAX = flags.DEFINE_integer("nconmax", None, "Maximum number of contacts.")
_NJMAX = flags.DEFINE_integer("njmax", None, "Maximum number of constraints per world.")
_NCCDMAX = flags.DEFINE_integer("nccdmax", None, "Maximum number of CCD contacts per world.")
_OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation: foo.bar = baz)", short_name="o")
_KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
_DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
@@ -134,27 +136,37 @@ def _main(argv: Sequence[str]) -> None:
else:
wp.config.quiet = flags.FLAGS["verbosity"].value < 1
wp.init()
if _CLEAR_KERNEL_CACHE.value:
wp.set_device(_DEVICE.value)
if _CLEAR_WARP_CACHE.value:
wp.clear_kernel_cache()
wp.clear_lto_cache()
# Clear CUDA compute cache for truly cold start JIT
compute_cache = epath.Path("~/.nv/ComputeCache").expanduser()
if compute_cache.exists():
shutil.rmtree(compute_cache)
compute_cache.mkdir()
with wp.ScopedDevice(_DEVICE.value):
m = mjw.put_model(mjm)
override_model(m, _OVERRIDE.value)
broadphase, filter = mjw.BroadphaseType(m.opt.broadphase).name, mjw.BroadphaseFilter(m.opt.broadphase_filter).name
solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name
integrator = mjw.IntegratorType(m.opt.integrator).name
iterations, ls_iterations = m.opt.iterations, m.opt.ls_iterations
ls_str = f"{'parallel' if m.opt.ls_parallel else 'iterative'} linesearch iterations: {ls_iterations}"
print(
f" nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.opt.is_sparse}\n"
f" broadphase: {broadphase} broadphase_filter: {filter}\n"
f" solver: {solver} cone: {cone} iterations: {iterations} {ls_str}\n"
f" integrator: {integrator} graph_conditional: {m.opt.graph_conditional}"
)
d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value)
print(f"Data\n nworld: {d.nworld} nconmax: {d.naconmax / d.nworld} njmax: {d.njmax}\n")
graph = _compile_step(m, d)
print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
override_model(mjm, _OVERRIDE.value)
m = mjw.put_model(mjm)
override_model(m, _OVERRIDE.value)
d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value, nccdmax=_NCCDMAX.value)
graph = _compile_step(m, d) if wp.get_device().is_cuda else None
if graph is None:
mjw.step(m, d) # warmup step
print("Running Warp unoptimized on CPU.")
broadphase, filter = mjw.BroadphaseType(m.opt.broadphase).name, mjw.BroadphaseFilter(m.opt.broadphase_filter).name
solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name
integrator = mjw.IntegratorType(m.opt.integrator).name
iterations, ls_iterations = m.opt.iterations, m.opt.ls_iterations
ls_str = f"{'parallel' if m.opt.ls_parallel else 'iterative'} linesearch iterations: {ls_iterations}"
print(
f" nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.is_sparse}\n"
f" broadphase: {broadphase} broadphase_filter: {filter}\n"
f" solver: {solver} cone: {cone} iterations: {iterations} {ls_str}\n"
f" integrator: {integrator} graph_conditional: {m.opt.graph_conditional}"
)
print(f"Data\n nworld: {d.nworld} nconmax: {int(d.naconmax / d.nworld)} njmax: {d.njmax}\n")
print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
opt = copy.copy(mjm.opt)
@@ -175,22 +187,19 @@ def _main(argv: Sequence[str]) -> None:
wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
# if the user changed an option in the MuJoCo Simulate UI, go ahead and recompile the step
# TODO: update memory tied to option max iterations
if mjm.opt != opt:
opt = copy.copy(mjm.opt)
m = mjw.put_model(mjm)
graph = _compile_step(m, d)
if _VIEWER_GLOBAL_STATE["running"]:
wp.capture_launch(graph)
wp.synchronize()
elif _VIEWER_GLOBAL_STATE["step_once"]:
graph = _compile_step(m, d) if wp.get_device().is_cuda else None
if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
_VIEWER_GLOBAL_STATE["step_once"] = False
wp.capture_launch(graph)
wp.synchronize()
if graph is None:
mjw.step(m, d)
else:
wp.capture_launch(graph)
wp.synchronize()
mjw.get_data_into(mjd, mjm, d)
viewer.sync()
mjx/mujoco/mjx/warp/collision_driver.py
+20 -34
View File
@@ -42,6 +42,7 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
)
@ffi.format_args_for_warp
def _collision_shim(
# Model
@@ -112,10 +113,8 @@ def _collision_shim(
opt__sdf_initpoints: int,
opt__sdf_iterations: int,
# Data
naccdmax: int,
naconmax: int,
collision_pair: wp.array(dtype=wp.vec2i),
collision_pairid: wp.array(dtype=wp.vec2i),
collision_worldid: wp.array(dtype=int),
geom_xmat: wp.array2d(dtype=wp.mat33),
geom_xpos: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
@@ -203,9 +202,6 @@ def _collision_shim(
_m.pair_solreffriction = pair_solreffriction
_m.plugin = plugin
_m.plugin_attr = plugin_attr
_d.collision_pair = collision_pair
_d.collision_pairid = collision_pairid
_d.collision_worldid = collision_worldid
_d.contact.dim = contact__dim
_d.contact.dist = contact__dist
_d.contact.frame = contact__frame
@@ -221,6 +217,7 @@ def _collision_shim(
_d.contact.worldid = contact__worldid
_d.geom_xmat = geom_xmat
_d.geom_xpos = geom_xpos
_d.naccdmax = naccdmax
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision
@@ -230,9 +227,6 @@ def _collision_shim(
def _collision_jax_impl(m: types.Model, d: types.Data):
output_dims = {
'collision_pair': d._impl.collision_pair.shape,
'collision_pairid': d._impl.collision_pairid.shape,
'collision_worldid': d._impl.collision_worldid.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape,
'contact__dim': d._impl.contact__dim.shape,
@@ -251,13 +245,10 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_collision_shim,
num_outputs=18,
num_outputs=15,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
'collision_pair',
'collision_pairid',
'collision_worldid',
'nacon',
'ncollision',
'contact__dim',
@@ -364,10 +355,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.opt.enableflags,
m.opt._impl.sdf_initpoints,
m.opt._impl.sdf_iterations,
d._impl.naccdmax,
d._impl.naconmax,
d._impl.collision_pair,
d._impl.collision_pairid,
d._impl.collision_worldid,
d.geom_xmat,
d.geom_xpos,
d._impl.nacon,
@@ -387,24 +376,21 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
d._impl.contact__worldid,
)
d = d.tree_replace({
'_impl.collision_pair': out[0],
'_impl.collision_pairid': out[1],
'_impl.collision_worldid': out[2],
'_impl.nacon': out[3],
'_impl.ncollision': out[4],
'_impl.contact__dim': out[5],
'_impl.contact__dist': out[6],
'_impl.contact__frame': out[7],
'_impl.contact__friction': out[8],
'_impl.contact__geom': out[9],
'_impl.contact__geomcollisionid': out[10],
'_impl.contact__includemargin': out[11],
'_impl.contact__pos': out[12],
'_impl.contact__solimp': out[13],
'_impl.contact__solref': out[14],
'_impl.contact__solreffriction': out[15],
'_impl.contact__type': out[16],
'_impl.contact__worldid': out[17],
'_impl.nacon': out[0],
'_impl.ncollision': out[1],
'_impl.contact__dim': out[2],
'_impl.contact__dist': out[3],
'_impl.contact__frame': out[4],
'_impl.contact__friction': out[5],
'_impl.contact__geom': out[6],
'_impl.contact__geomcollisionid': out[7],
'_impl.contact__includemargin': out[8],
'_impl.contact__pos': out[9],
'_impl.contact__solimp': out[10],
'_impl.contact__solref': out[11],
'_impl.contact__solreffriction': out[12],
'_impl.contact__type': out[13],
'_impl.contact__worldid': out[14],
})
return d
mjx/mujoco/mjx/warp/forward.py
+252 -535
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/smooth.py
+13 -3
View File
@@ -42,10 +42,13 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
)
@ffi.format_args_for_warp
def _kinematics_shim(
# Model
nworld: int,
body_branch_start: wp.array(dtype=int),
body_branches: wp.array(dtype=int),
body_ipos: wp.array2d(dtype=wp.vec3),
body_iquat: wp.array2d(dtype=wp.quat),
body_jntadr: wp.array(dtype=int),
@@ -55,7 +58,6 @@ def _kinematics_shim(
body_pos: wp.array2d(dtype=wp.vec3),
body_quat: wp.array2d(dtype=wp.quat),
body_rootid: wp.array(dtype=int),
body_tree: tuple[wp.array(dtype=int), ...],
body_weldid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_pos: wp.array2d(dtype=wp.vec3),
@@ -64,6 +66,8 @@ def _kinematics_shim(
jnt_pos: wp.array2d(dtype=wp.vec3),
jnt_qposadr: wp.array(dtype=int),
jnt_type: wp.array(dtype=int),
nbody: int,
nbranch: int,
ngeom: int,
nsite: int,
qpos0: wp.array2d(dtype=float),
@@ -90,6 +94,8 @@ def _kinematics_shim(
_m.opt = _o
_d.efc = _e
_d.contact = _c
_m.body_branch_start = body_branch_start
_m.body_branches = body_branches
_m.body_ipos = body_ipos
_m.body_iquat = body_iquat
_m.body_jntadr = body_jntadr
@@ -99,7 +105,6 @@ def _kinematics_shim(
_m.body_pos = body_pos
_m.body_quat = body_quat
_m.body_rootid = body_rootid
_m.body_tree = body_tree
_m.body_weldid = body_weldid
_m.geom_bodyid = geom_bodyid
_m.geom_pos = geom_pos
@@ -108,6 +113,8 @@ def _kinematics_shim(
_m.jnt_pos = jnt_pos
_m.jnt_qposadr = jnt_qposadr
_m.jnt_type = jnt_type
_m.nbody = nbody
_m.nbranch = nbranch
_m.ngeom = ngeom
_m.nsite = nsite
_m.qpos0 = qpos0
@@ -208,6 +215,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
)
out = jf(
d.qpos.shape[0],
m._impl.body_branch_start,
m._impl.body_branches,
m.body_ipos,
m.body_iquat,
m.body_jntadr,
@@ -217,7 +226,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
m.body_pos,
m.body_quat,
m.body_rootid,
m._impl.body_tree,
m.body_weldid,
m.geom_bodyid,
m.geom_pos,
@@ -226,6 +234,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
m.jnt_pos,
m.jnt_qposadr,
m.jnt_type,
m.nbody,
m._impl.nbranch,
m.ngeom,
m.nsite,
m.qpos0,
mjx/mujoco/mjx/warp/types.py
+70 -121
View File
@@ -72,7 +72,7 @@ class BlockDim:
energy_vel_kinetic: int
euler_dense: int
linesearch_iterative: int
mul_m_dense: int
qderiv_actuator_dense: int
ray: int
segmented_sort: int
tendon_velocity: int
@@ -93,7 +93,7 @@ class BlockDim:
class StatisticWarp(PyTreeNode):
"""Derived fields from Statistic."""
meaninertia: float
meaninertia: jax.Array
class OptionWarp(PyTreeNode):
"""Derived fields from Option."""
@@ -104,9 +104,7 @@ class OptionWarp(PyTreeNode):
contact_sensor_maxmatch: int
graph_conditional: bool
graph_mode: GraphMode
has_fluid: bool
impratio_invsqrt: jax.Array
is_sparse: bool
ls_parallel: bool
ls_parallel_min_step: float
run_collision_detection: bool
@@ -120,8 +118,11 @@ class ModelWarp(PyTreeNode):
M_rownnz: np.ndarray
actuator_trntype_body_adr: np.ndarray
block_dim: BlockDim
body_branch_start: np.ndarray
body_branches: np.ndarray
body_fluid_ellipsoid: np.ndarray
body_tree: Tuple[np.ndarray, ...]
cam_projection: np.ndarray
collision_sensor_adr: np.ndarray
dof_tri_col: np.ndarray
dof_tri_row: np.ndarray
@@ -146,11 +147,16 @@ class ModelWarp(PyTreeNode):
flex_vertadr: np.ndarray
flex_vertbodyid: np.ndarray
flex_vertnum: np.ndarray
flexedge_J_colind: np.ndarray
flexedge_J_rowadr: np.ndarray
flexedge_J_rownnz: np.ndarray
flexedge_invweight0: np.ndarray
flexedge_length0: np.ndarray
geom_pair_type_count: Tuple[int, ...]
geom_plugin_index: np.ndarray
has_fluid: bool
has_sdf_geom: bool
is_sparse: bool
jnt_limited_ball_adr: np.ndarray
jnt_limited_slide_hinge_adr: np.ndarray
light_active: jax.Array
@@ -169,6 +175,7 @@ class ModelWarp(PyTreeNode):
mesh_polyvertnum: np.ndarray
mocap_bodyid: np.ndarray
nacttrnbody: int
nbranch: int
nflex: int
nflexedge: int
nflexelem: int
@@ -185,6 +192,7 @@ class ModelWarp(PyTreeNode):
nsensorcollision: int
nsensorcontact: int
nsensortaxel: int
ntree: int
nv_pad: int
nxn_geom_pair: np.ndarray
nxn_geom_pair_filtered: np.ndarray
@@ -198,9 +206,9 @@ class ModelWarp(PyTreeNode):
qLD_updates: Tuple[np.ndarray, ...]
qM_fullm_i: np.ndarray
qM_fullm_j: np.ndarray
qM_madr_ij: np.ndarray
qM_mulm_i: np.ndarray
qM_mulm_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
@@ -228,6 +236,9 @@ class ModelWarp(PyTreeNode):
tendon_jnt_adr: np.ndarray
tendon_limited_adr: np.ndarray
tendon_site_pair_adr: np.ndarray
tree_bodynum: np.ndarray
tree_dofadr: np.ndarray
tree_dofnum: np.ndarray
wrap_geom_adr: np.ndarray
wrap_jnt_adr: np.ndarray
wrap_pulley_scale: np.ndarray
@@ -242,9 +253,6 @@ class DataWarp(PyTreeNode):
cfrc_ext: jax.Array
cfrc_int: jax.Array
cinert: jax.Array
collision_pair: jax.Array
collision_pairid: jax.Array
collision_worldid: jax.Array
contact__dim: jax.Array
contact__dist: jax.Array
contact__efc_address: jax.Array
@@ -262,31 +270,13 @@ class DataWarp(PyTreeNode):
crb: jax.Array
efc__D: jax.Array
efc__J: jax.Array
efc__Jaref: jax.Array
efc__Ma: jax.Array
efc__Mgrad: jax.Array
efc__alpha: jax.Array
efc__aref: jax.Array
efc__beta: jax.Array
efc__cost: jax.Array
efc__done: jax.Array
efc__force: jax.Array
efc__frictionloss: jax.Array
efc__gauss: jax.Array
efc__grad: jax.Array
efc__grad_dot: jax.Array
efc__id: jax.Array
efc__jv: jax.Array
efc__margin: jax.Array
efc__mv: jax.Array
efc__pos: jax.Array
efc__prev_Mgrad: jax.Array
efc__prev_cost: jax.Array
efc__prev_grad: jax.Array
efc__quad: jax.Array
efc__quad_gauss: jax.Array
efc__search: jax.Array
efc__search_dot: jax.Array
efc__state: jax.Array
efc__type: jax.Array
efc__vel: jax.Array
@@ -297,20 +287,15 @@ class DataWarp(PyTreeNode):
flexvert_xpos: jax.Array
light_xdir: jax.Array
light_xpos: jax.Array
naccdmax: int
nacon: jax.Array
naconmax: int
ncollision: jax.Array
ne: jax.Array
ne_connect: jax.Array
ne_flex: jax.Array
ne_jnt: jax.Array
ne_ten: jax.Array
ne_weld: jax.Array
nefc: jax.Array
nf: jax.Array
njmax: int
nl: jax.Array
nsolving: jax.Array
nworld: int
qLD: jax.Array
qLDiagInv: jax.Array
@@ -319,7 +304,6 @@ class DataWarp(PyTreeNode):
qfrc_spring: jax.Array
solver_niter: jax.Array
subtree_angmom: jax.Array
subtree_bodyvel: jax.Array
subtree_linvel: jax.Array
ten_J: jax.Array
ten_velocity: jax.Array
@@ -329,9 +313,6 @@ class DataWarp(PyTreeNode):
wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
'collision_pair',
'collision_pairid',
'collision_worldid',
'contact__dim',
'contact__dist',
'contact__efc_address',
@@ -346,11 +327,11 @@ DATA_NON_VMAP = {
'contact__solreffriction',
'contact__type',
'contact__worldid',
'naccdmax',
'nacon',
'naconmax',
'ncollision',
'njmax',
'nsolving',
'nworld',
}
@@ -394,9 +375,6 @@ _NDIM = {
'cfrc_ext': 3,
'cfrc_int': 3,
'cinert': 3,
'collision_pair': 2,
'collision_pairid': 2,
'collision_worldid': 1,
'contact__dim': 1,
'contact__dist': 1,
'contact__efc_address': 2,
@@ -416,31 +394,13 @@ _NDIM = {
'cvel': 3,
'efc__D': 2,
'efc__J': 3,
'efc__Jaref': 2,
'efc__Ma': 2,
'efc__Mgrad': 2,
'efc__alpha': 1,
'efc__aref': 2,
'efc__beta': 1,
'efc__cost': 1,
'efc__done': 1,
'efc__force': 2,
'efc__frictionloss': 2,
'efc__gauss': 1,
'efc__grad': 2,
'efc__grad_dot': 1,
'efc__id': 2,
'efc__jv': 2,
'efc__margin': 2,
'efc__mv': 2,
'efc__pos': 2,
'efc__prev_Mgrad': 2,
'efc__prev_cost': 1,
'efc__prev_grad': 2,
'efc__quad': 3,
'efc__quad_gauss': 2,
'efc__search': 2,
'efc__search_dot': 1,
'efc__state': 2,
'efc__type': 2,
'efc__vel': 2,
@@ -456,20 +416,15 @@ _NDIM = {
'light_xpos': 3,
'mocap_pos': 3,
'mocap_quat': 3,
'naccdmax': 0,
'nacon': 1,
'naconmax': 0,
'ncollision': 1,
'ne': 1,
'ne_connect': 1,
'ne_flex': 1,
'ne_jnt': 1,
'ne_ten': 1,
'ne_weld': 1,
'nefc': 1,
'nf': 1,
'njmax': 0,
'nl': 1,
'nsolving': 1,
'nworld': 0,
'qLD': 3,
'qLDiagInv': 2,
@@ -495,7 +450,6 @@ _NDIM = {
'site_xpos': 3,
'solver_niter': 1,
'subtree_angmom': 3,
'subtree_bodyvel': 3,
'subtree_com': 3,
'subtree_linvel': 3,
'ten_J': 3,
@@ -549,7 +503,7 @@ _NDIM = {
'block_dim__energy_vel_kinetic': 0,
'block_dim__euler_dense': 0,
'block_dim__linesearch_iterative': 0,
'block_dim__mul_m_dense': 0,
'block_dim__qderiv_actuator_dense': 0,
'block_dim__ray': 0,
'block_dim__segmented_sort': 0,
'block_dim__tendon_velocity': 0,
@@ -557,6 +511,8 @@ _NDIM = {
'block_dim__update_gradient_JTDAJ_sparse': 0,
'block_dim__update_gradient_cholesky': 0,
'block_dim__update_gradient_cholesky_blocked': 0,
'body_branch_start': 1,
'body_branches': 1,
'body_conaffinity': 1,
'body_contype': 1,
'body_dofadr': 1,
@@ -579,15 +535,17 @@ _NDIM = {
'body_rootid': 1,
'body_subtreemass': 2,
'body_tree': -1,
'body_treeid': 1,
'body_weldid': 1,
'cam_bodyid': 1,
'cam_fovy': 1,
'cam_intrinsic': 2,
'cam_fovy': 2,
'cam_intrinsic': 3,
'cam_mat0': 4,
'cam_mode': 1,
'cam_pos': 3,
'cam_pos0': 3,
'cam_poscom0': 3,
'cam_projection': 1,
'cam_quat': 3,
'cam_resolution': 2,
'cam_sensorsize': 2,
@@ -603,6 +561,7 @@ _NDIM = {
'dof_parentid': 1,
'dof_solimp': 3,
'dof_solref': 3,
'dof_treeid': 1,
'dof_tri_col': 1,
'dof_tri_row': 1,
'eq_active0': 1,
@@ -635,6 +594,9 @@ _NDIM = {
'flex_vertadr': 1,
'flex_vertbodyid': 1,
'flex_vertnum': 1,
'flexedge_J_colind': 1,
'flexedge_J_rowadr': 1,
'flexedge_J_rownnz': 1,
'flexedge_invweight0': 1,
'flexedge_length0': 1,
'geom_aabb': 4,
@@ -661,12 +623,14 @@ _NDIM = {
'geom_solmix': 2,
'geom_solref': 3,
'geom_type': 1,
'has_fluid': 0,
'has_sdf_geom': 0,
'hfield_adr': 1,
'hfield_data': 1,
'hfield_ncol': 1,
'hfield_nrow': 1,
'hfield_size': 2,
'is_sparse': 0,
'jnt_actfrclimited': 1,
'jnt_actfrcrange': 3,
'jnt_actgravcomp': 1,
@@ -697,7 +661,7 @@ _NDIM = {
'light_type': 2,
'mapM2M': 1,
'mat_rgba': 3,
'mat_texid': 2,
'mat_texid': 3,
'mat_texrepeat': 3,
'mesh_face': 2,
'mesh_faceadr': 1,
@@ -724,6 +688,7 @@ _NDIM = {
'na': 0,
'nacttrnbody': 0,
'nbody': 0,
'nbranch': 0,
'ncam': 0,
'neq': 0,
'nexclude': 0,
@@ -765,6 +730,7 @@ _NDIM = {
'nsensortaxel': 0,
'nsite': 0,
'ntendon': 0,
'ntree': 0,
'nu': 0,
'nv': 0,
'nv_pad': 0,
@@ -785,10 +751,8 @@ _NDIM = {
'opt__enableflags': 0,
'opt__graph_conditional': 0,
'opt__gravity': 2,
'opt__has_fluid': 0,
'opt__impratio_invsqrt': 1,
'opt__integrator': 0,
'opt__is_sparse': 0,
'opt__iterations': 0,
'opt__ls_iterations': 0,
'opt__ls_parallel': 0,
@@ -817,9 +781,9 @@ _NDIM = {
'qLD_updates': -1,
'qM_fullm_i': 1,
'qM_fullm_j': 1,
'qM_madr_ij': 1,
'qM_mulm_i': 1,
'qM_mulm_j': 1,
'qM_mulm_col': 1,
'qM_mulm_madr': 1,
'qM_mulm_rowadr': 1,
'qM_tiles': -1,
'qpos0': 2,
'qpos_spring': 2,
@@ -856,7 +820,7 @@ _NDIM = {
'site_quat': 3,
'site_size': 2,
'site_type': 1,
'stat__meaninertia': 0,
'stat__meaninertia': 1,
'taxel_sensorid': 1,
'taxel_vertadr': 1,
'ten_wrapadr_site': 1,
@@ -883,6 +847,9 @@ _NDIM = {
'tendon_solref_fri': 3,
'tendon_solref_lim': 3,
'tendon_stiffness': 2,
'tree_bodynum': 1,
'tree_dofadr': 1,
'tree_dofnum': 1,
'wrap_geom_adr': 1,
'wrap_jnt_adr': 1,
'wrap_objid': 1,
@@ -902,10 +869,8 @@ _NDIM = {
'enableflags': 0,
'graph_conditional': 0,
'gravity': 2,
'has_fluid': 0,
'impratio_invsqrt': 1,
'integrator': 0,
'is_sparse': 0,
'iterations': 0,
'ls_iterations': 0,
'ls_parallel': 0,
@@ -921,7 +886,7 @@ _NDIM = {
'viscosity': 1,
'wind': 2,
},
'Statistic': {'meaninertia': 0},
'Statistic': {'meaninertia': 1},
}
_BATCH_DIM = {
'Data': {
@@ -939,9 +904,6 @@ _BATCH_DIM = {
'cfrc_ext': True,
'cfrc_int': True,
'cinert': True,
'collision_pair': False,
'collision_pairid': False,
'collision_worldid': False,
'contact__dim': False,
'contact__dist': False,
'contact__efc_address': False,
@@ -961,31 +923,13 @@ _BATCH_DIM = {
'cvel': True,
'efc__D': True,
'efc__J': True,
'efc__Jaref': True,
'efc__Ma': True,
'efc__Mgrad': True,
'efc__alpha': True,
'efc__aref': True,
'efc__beta': True,
'efc__cost': True,
'efc__done': True,
'efc__force': True,
'efc__frictionloss': True,
'efc__gauss': True,
'efc__grad': True,
'efc__grad_dot': True,
'efc__id': True,
'efc__jv': True,
'efc__margin': True,
'efc__mv': True,
'efc__pos': True,
'efc__prev_Mgrad': True,
'efc__prev_cost': True,
'efc__prev_grad': True,
'efc__quad': True,
'efc__quad_gauss': True,
'efc__search': True,
'efc__search_dot': True,
'efc__state': True,
'efc__type': True,
'efc__vel': True,
@@ -1001,20 +945,15 @@ _BATCH_DIM = {
'light_xpos': True,
'mocap_pos': True,
'mocap_quat': True,
'naccdmax': False,
'nacon': False,
'naconmax': False,
'ncollision': False,
'ne': True,
'ne_connect': True,
'ne_flex': True,
'ne_jnt': True,
'ne_ten': True,
'ne_weld': True,
'nefc': True,
'nf': True,
'njmax': False,
'nl': True,
'nsolving': False,
'nworld': False,
'qLD': True,
'qLDiagInv': True,
@@ -1040,7 +979,6 @@ _BATCH_DIM = {
'site_xpos': True,
'solver_niter': True,
'subtree_angmom': True,
'subtree_bodyvel': True,
'subtree_com': True,
'subtree_linvel': True,
'ten_J': True,
@@ -1094,7 +1032,7 @@ _BATCH_DIM = {
'block_dim__energy_vel_kinetic': False,
'block_dim__euler_dense': False,
'block_dim__linesearch_iterative': False,
'block_dim__mul_m_dense': False,
'block_dim__qderiv_actuator_dense': False,
'block_dim__ray': False,
'block_dim__segmented_sort': False,
'block_dim__tendon_velocity': False,
@@ -1102,6 +1040,8 @@ _BATCH_DIM = {
'block_dim__update_gradient_JTDAJ_sparse': False,
'block_dim__update_gradient_cholesky': False,
'block_dim__update_gradient_cholesky_blocked': False,
'body_branch_start': False,
'body_branches': False,
'body_conaffinity': False,
'body_contype': False,
'body_dofadr': False,
@@ -1124,15 +1064,17 @@ _BATCH_DIM = {
'body_rootid': False,
'body_subtreemass': True,
'body_tree': False,
'body_treeid': False,
'body_weldid': False,
'cam_bodyid': False,
'cam_fovy': False,
'cam_intrinsic': False,
'cam_fovy': True,
'cam_intrinsic': True,
'cam_mat0': True,
'cam_mode': False,
'cam_pos': True,
'cam_pos0': True,
'cam_poscom0': True,
'cam_projection': False,
'cam_quat': True,
'cam_resolution': False,
'cam_sensorsize': False,
@@ -1148,6 +1090,7 @@ _BATCH_DIM = {
'dof_parentid': False,
'dof_solimp': True,
'dof_solref': True,
'dof_treeid': False,
'dof_tri_col': False,
'dof_tri_row': False,
'eq_active0': False,
@@ -1180,6 +1123,9 @@ _BATCH_DIM = {
'flex_vertadr': False,
'flex_vertbodyid': False,
'flex_vertnum': False,
'flexedge_J_colind': False,
'flexedge_J_rowadr': False,
'flexedge_J_rownnz': False,
'flexedge_invweight0': False,
'flexedge_length0': False,
'geom_aabb': True,
@@ -1206,12 +1152,14 @@ _BATCH_DIM = {
'geom_solmix': True,
'geom_solref': True,
'geom_type': False,
'has_fluid': False,
'has_sdf_geom': False,
'hfield_adr': False,
'hfield_data': False,
'hfield_ncol': False,
'hfield_nrow': False,
'hfield_size': False,
'is_sparse': False,
'jnt_actfrclimited': False,
'jnt_actfrcrange': True,
'jnt_actgravcomp': False,
@@ -1242,7 +1190,7 @@ _BATCH_DIM = {
'light_type': True,
'mapM2M': False,
'mat_rgba': True,
'mat_texid': False,
'mat_texid': True,
'mat_texrepeat': True,
'mesh_face': False,
'mesh_faceadr': False,
@@ -1269,6 +1217,7 @@ _BATCH_DIM = {
'na': False,
'nacttrnbody': False,
'nbody': False,
'nbranch': False,
'ncam': False,
'neq': False,
'nexclude': False,
@@ -1310,6 +1259,7 @@ _BATCH_DIM = {
'nsensortaxel': False,
'nsite': False,
'ntendon': False,
'ntree': False,
'nu': False,
'nv': False,
'nv_pad': False,
@@ -1330,10 +1280,8 @@ _BATCH_DIM = {
'opt__enableflags': False,
'opt__graph_conditional': False,
'opt__gravity': True,
'opt__has_fluid': False,
'opt__impratio_invsqrt': True,
'opt__integrator': False,
'opt__is_sparse': False,
'opt__iterations': False,
'opt__ls_iterations': False,
'opt__ls_parallel': False,
@@ -1362,9 +1310,9 @@ _BATCH_DIM = {
'qLD_updates': False,
'qM_fullm_i': False,
'qM_fullm_j': False,
'qM_madr_ij': False,
'qM_mulm_i': False,
'qM_mulm_j': False,
'qM_mulm_col': False,
'qM_mulm_madr': False,
'qM_mulm_rowadr': False,
'qM_tiles': False,
'qpos0': True,
'qpos_spring': True,
@@ -1401,7 +1349,7 @@ _BATCH_DIM = {
'site_quat': True,
'site_size': False,
'site_type': False,
'stat__meaninertia': False,
'stat__meaninertia': True,
'taxel_sensorid': False,
'taxel_vertadr': False,
'ten_wrapadr_site': False,
@@ -1428,6 +1376,9 @@ _BATCH_DIM = {
'tendon_solref_fri': True,
'tendon_solref_lim': True,
'tendon_stiffness': True,
'tree_bodynum': False,
'tree_dofadr': False,
'tree_dofnum': False,
'wrap_geom_adr': False,
'wrap_jnt_adr': False,
'wrap_objid': False,
@@ -1447,10 +1398,8 @@ _BATCH_DIM = {
'enableflags': False,
'graph_conditional': False,
'gravity': True,
'has_fluid': False,
'impratio_invsqrt': True,
'integrator': False,
'is_sparse': False,
'iterations': False,
'ls_iterations': False,
'ls_parallel': False,
@@ -1466,5 +1415,5 @@ _BATCH_DIM = {
'viscosity': True,
'wind': True,
},
'Statistic': {'meaninertia': False},
'Statistic': {'meaninertia': True},
}