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: 799614257
Change-Id: Ie05d4dd7c45308ed5e8010c3eca4b68f552335ea
This commit is contained in:
Baruch Tabanpour
2025-08-26 10:16:17 -07:00
committed by Copybara-Service
parent 1ac51cf2a5
commit 9eaa31af2f
20 changed files with 701 additions and 1552 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+22 -12
View File
@@ -104,7 +104,7 @@ def _max_contacts_height_field(
@cache_kernel
def ccd_kernel_builder(
default_gjk: bool,
legacy_gjk: bool,
geomtype1: int,
geomtype2: int,
gjk_iterations: int,
@@ -286,10 +286,16 @@ def ccd_kernel_builder(
hftri_index,
)
points = mat3c()
# TODO(kbayes): remove legacy GJK once multicontact can be enabled
if default_gjk:
if wp.static(legacy_gjk):
# find prism center for height field
if geomtype1 == int(GeomType.HFIELD.value):
x1 = wp.vec3(0.0, 0.0, 0.0)
for i in range(6):
x1 += hfield_prism_vertex(geom1.hfprism, i)
x1 = geom1.pos + geom1.rot @ (x1 / 6.0)
geom1.pos = x1
simplex, normal = gjk_legacy(
gjk_iterations,
geom1,
@@ -303,24 +309,28 @@ def ccd_kernel_builder(
)
dist = -depth
if (dist - margin) >= 0.0 or depth != depth:
if dist >= 0.0 or depth < -depth_extension:
count = 0
return
sphere = int(GeomType.SPHERE.value)
ellipsoid = int(GeomType.ELLIPSOID.value)
if geom_type[g1] == sphere or geom_type[g1] == ellipsoid or geom_type[g2] == sphere or geom_type[g2] == ellipsoid:
count, points = multicontact_legacy(geom1, geom2, geomtype1, geomtype2, depth_extension, depth, normal, 1, 2, 1.0e-5)
else:
count, points = multicontact_legacy(geom1, geom2, geomtype1, geomtype2, depth_extension, depth, normal, 4, 8, 1.0e-3)
count, points = multicontact_legacy(geom1, geom2, geomtype1, geomtype2, depth_extension, depth, normal, 4, 8, 1.0e-1)
frame = make_frame(normal)
else:
points = mat3c()
x1 = geom1.pos
x2 = geom2.pos
# find prism center for height field
if geomtype1 == int(GeomType.HFIELD.value):
x1 = wp.vec3(0.0, 0.0, 0.0)
x1_ = wp.vec3(0.0, 0.0, 0.0)
for i in range(6):
x1 += hfield_prism_vertex(geom1.hfprism, i)
x1 = x1 / 6.0
x1_ += hfield_prism_vertex(geom1.hfprism, i)
x1 += geom1.rot @ (x1_ / 6.0)
dist, count, witness1, witness2 = ccd(
False,
@@ -411,13 +421,13 @@ def convex_narrowphase(m: Model, d: Data):
if m.geom_pair_type_count[upper_trid_index(len(GeomType), geom_pair[0], geom_pair[1])]:
wp.launch(
ccd_kernel_builder(
False,
m.opt.legacy_gjk,
geom_pair[0],
geom_pair[1],
m.opt.gjk_iterations,
m.opt.epa_iterations,
False,
0.1,
True,
1e9,
),
dim=d.nconmax,
inputs=[
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+1 -1
View File
@@ -113,7 +113,7 @@ def _support_margin(geom: Geom, geomtype: int, dir: wp.vec3):
local_dir = wp.transpose(geom.rot) @ dir
res = wp.vec3()
res[2] = wp.where(local_dir[2] >= 0, geom.size[1], -geom.size[1])
sp.point = res
sp.point = geom.rot @ res + geom.pos
return sp
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_legacy.py
Vendored
+10 -8
View File
@@ -19,6 +19,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_hfield import hfield_pris
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
from mujoco.mjx.third_party.mujoco_warp._src.math import gjk_normalize
from mujoco.mjx.third_party.mujoco_warp._src.math import orthonormal
from mujoco.mjx.third_party.mujoco_warp._src.math import orthonormal_to_z
from mujoco.mjx.third_party.mujoco_warp._src.support import all_same
from mujoco.mjx.third_party.mujoco_warp._src.support import any_different
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
@@ -203,8 +204,8 @@ def gjk_legacy(
depth = dist_min
normal = dir_n
sd = simplex0 - simplex1
dir = orthonormal(sd)
sd = wp.normalize(simplex0 - simplex1)
dir = orthonormal_to_z(sd)
dist_max, simplex3 = _gjk_support(geom1, geom2, geomtype1, geomtype2, dir)
@@ -304,14 +305,15 @@ def epa_legacy(
normal: wp.vec3,
):
# get the support, if depth < 0: objects do not intersect
depth, _ = _gjk_support(geom1, geom2, geomtype1, geomtype2, normal)
depth, simplex0 = _gjk_support(geom1, geom2, geomtype1, geomtype2, normal)
simplex[0] = simplex0
if depth < -depth_extension:
# Objects are not intersecting, and we do not obtain the closest points as
# specified by depth_extension.
return wp.nan, wp.vec3(wp.nan, wp.nan, wp.nan)
return FLOAT_MAX, wp.vec3(wp.nan, wp.nan, wp.nan)
if wp.static(epa_exact_neg_distance):
if epa_exact_neg_distance:
# Check closest points to all edges of the simplex, rather than just the
# face normals. This gives the exact depth/normal for the non-intersecting
# case.
@@ -364,7 +366,7 @@ def epa_legacy(
# This is a hack to reduce compile time
count = int(4)
it = int(0)
for _ in range(wp.static(epa_iterations)):
for _ in range(epa_iterations):
it += count
count = wp.min(count * 3, EPS_BEST_COUNT)
@@ -391,7 +393,7 @@ def epa_legacy(
# iterate over edges and get distance using support point
for j in range(3):
if wp.static(epa_exact_neg_distance):
if epa_exact_neg_distance:
# obtain closest point between new triangle edge and origin
tqj = tris[ti + j]
@@ -663,7 +665,7 @@ def multicontact_legacy(
m2 = v2[j]
dd = (m1[0] - m2[0]) * (m1[0] - m2[0]) + (m1[1] - m2[1]) * (m1[1] - m2[1])
if i != 0 and j != 0 or dd < minDist:
if (i == 0 and j == 0) or (dd < minDist):
minDist = dd
var_rx = ((m1[0] + m2[0]) * dir + (m1[1] + m2[1]) * dir2 + (m1[2] + m2[2]) * normal) * 0.5
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_hfield.py
Vendored
+166 -162
View File
@@ -17,74 +17,49 @@ from typing import Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@wp.func
def _hfield_overlap_range(
# Model:
geom_dataid: wp.array(dtype=int),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
hfield_nrow: wp.array(dtype=int),
hfield_ncol: wp.array(dtype=int),
hfield_size: wp.array(dtype=wp.vec4),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
def _hfield_subgrid(
# In:
hfieldid: int,
geomid: int,
worldid: int,
nrow: int,
ncol: int,
size: wp.vec4,
xmax: float,
xmin: float,
ymax: float,
ymin: float,
) -> Tuple[int, int, int, int]:
"""Returns min/max grid coordinates of height field cells overlapped by a geom's bounds.
"""Returns height field subgrid that overlaps with geom AABB.
Args:
geom_dataid: Array of geom data IDs
geom_rbound: Array of geom bounding radii
geom_margin: Array of geom margins
hfield_nrow: Array of heightfield rows
hfield_ncol: Array of heightfield columns
hfield_size: Array of heightfield sizes
geom_xpos_in: Array of geom positions
geom_xmat_in: Array of geom orientation matrices
hfieldid: Index of the height field geom
geomid: Index of the other geom
worldid: Current world index
nrow: height field number of rows
ncol: height field number of columns
size: height field size
xmax: geom maximum x position
xmin: geom minimum x position
ymax: geom maximum y position
ymin: geom minimum y position
Returns:
min_i, min_j, max_i, max_j: Grid coordinate bounds
grid coordinate bounds
"""
# get height field dimensions
dataid = geom_dataid[hfieldid]
nrow = hfield_nrow[dataid]
ncol = hfield_ncol[dataid]
size = hfield_size[dataid] # (x, y, z_top, z_bottom)
# get positions and transforms
hf_pos = geom_xpos_in[worldid, hfieldid]
hf_mat = geom_xmat_in[worldid, hfieldid]
geom_pos = geom_xpos_in[worldid, geomid]
# grid resolution
x_scale = 0.5 * float(ncol - 1) / size[0]
y_scale = 0.5 * float(nrow - 1) / size[1]
# transform geom_pos to height field local space
local_pos = wp.transpose(hf_mat) @ (geom_pos - hf_pos)
# subgrid
cmin = wp.max(0, int(wp.floor((xmin + size[0]) * x_scale)))
cmax = wp.min(ncol - 1, int(wp.ceil((xmax + size[0]) * x_scale)))
rmin = wp.max(0, int(wp.floor((ymin + size[1]) * y_scale)))
rmax = wp.min(nrow - 1, int(wp.ceil((ymax + size[1]) * y_scale)))
# get bounding radius of other geometry (including margin)
bound_radius = geom_rbound[worldid, geomid] + geom_margin[worldid, geomid]
# calculate grid resolution
x_scale = 2.0 * size[0] / float(ncol - 1)
y_scale = 2.0 * size[1] / float(nrow - 1)
# calculate min/max grid coordinates that could contain the object
min_i = wp.max(0, int((local_pos[0] - bound_radius + size[0]) / x_scale))
max_i = wp.min(ncol - 2, int((local_pos[0] + bound_radius + size[0]) / x_scale) + 1)
min_j = wp.max(0, int((local_pos[1] - bound_radius + size[1]) / y_scale))
max_j = wp.min(nrow - 2, int((local_pos[1] + bound_radius + size[1]) / y_scale) + 1)
return min_i, min_j, max_i, max_j
return cmin, rmin, cmax, rmax
@wp.func
@@ -100,20 +75,20 @@ def hfield_triangle_prism(
hfieldid: int,
hftri_index: int,
) -> wp.mat33:
"""Returns the vertices of a triangular prism for a heightfield triangle.
"""Returns triangular prism vertex information in compressed representation.
Args:
geom_dataid: Array of geometry data IDs
hfield_adr: Array of heightfield addresses
hfield_nrow: Array of heightfield rows
hfield_ncol: Array of heightfield columns
hfield_size: Array of heightfield sizes
hfield_data: Array of heightfield data
hfieldid: Index of the height field geometry
hftri_index: Index of the triangle in the heightfield
geom_dataid: geom data ids
hfield_adr: address for height field
hfield_nrow: height field number of rows
hfield_ncol: height field number of columns
hfield_size: height field sizes
hfield_data: height field data
hfieldid: height field geom id
hftri_index: height field triangle index
Returns:
3x3 matrix containing the vertices of the triangular prism
triangular prism vertex information (compressed)
"""
# https://mujoco.readthedocs.io/en/stable/XMLreference.html#asset-hfield
@@ -160,21 +135,14 @@ def hfield_triangle_prism(
z10 = z10 * z_top
z11 = z11 * z_top
# set bottom z-value
z_bottom = -size[3]
x2 = wp.where(hftri_index % 2, 1.0, 0.0)
y2 = wp.where(hftri_index % 2, z10, z01)
z22 = -size[3]
# compress 6 prism vertices into 3x3 matrix, see hfield_prism_vertex for details
return wp.mat33(
x0,
y0,
z00,
x1,
y1,
z11,
wp.where(hftri_index % 2, 1.0, 0.0),
wp.where(hftri_index % 2, z10, z01),
z_bottom,
)
return wp.mat33(x0, y0, z00,
x1, y1, z11,
x2, y2, z22) # fmt: off
@wp.func
@@ -194,10 +162,10 @@ def hfield_prism_vertex(prism: wp.mat33, vert_index: int) -> wp.vec3:
Args:
prism: 3x3 compressed representation of a triangular prism
vert_index: Index of vertex to extract (0-5)
vert_index: index of vertex to extract (0-5)
Returns:
The 3D coordinates of the requested vertex
3D coordinates of the requested vertex
"""
if vert_index == 0 or vert_index == 1:
return prism[vert_index] # first two vertices stored directly
@@ -223,6 +191,7 @@ def _hfield_midphase(
# Model:
geom_type: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_aabb: wp.array2d(dtype=wp.vec3),
geom_rbound: wp.array2d(dtype=float),
geom_margin: wp.array2d(dtype=float),
hfield_nrow: wp.array(dtype=int),
@@ -250,130 +219,171 @@ def _hfield_midphase(
one for each potentially colliding triangle.
Args:
geom_type: Array of geometry types
geom_dataid: Array of geometry data IDs
geom_rbound: Array of geometry bounding radii
geom_margin: Array of geometry margins
hfield_nrow: Array of heightfield rows
hfield_ncol: Array of heightfield columns
hfield_size: Array of heightfield sizes
nconmax_in: Max number of collisions
geom_xpos_in: Array of geometry positions
geom_xmat_in: Array of geometry orientation matrices
collision_pair_in: Array of collision pairs
collision_hftri_index_in: Array of heightfield triangle indices, -1 for heightfield
pairs
collision_pairid_in: Array of collision pair IDs
collision_worldid_in: Array of collision world IDs
collision_pair_out: Output array of collision pairs
collision_hftri_index_out: Output array of heightfield triangle indices
collision_pairid_out: Output array of collision pair IDs
collision_worldid_out: Output array of collision world IDs
ncollision_out: Output counter for number of collisions
geom_type: geom type
geom_dataid: geom data id
geom_rbound: geom bounding sphere radius
geom_margin: geom margin
hfield_nrow: height field number of rows
hfield_ncol: height field number of columns
hfield_size: height field size
nconmax_in: maximum number of contacts
geom_xpos_in: geom position
geom_xmat_in: geom orientation
collision_pair_in: collision pair
collision_hftri_index_in: triangle indices, -1 for height field pair
collision_pairid_in: collision pair id from broadphase
collision_worldid_in: collision world id from broadphase
collision_pair_out: collision pair from midphase
collision_hftri_index_out: triangle indices from midphase
collision_pairid_out: collision pair id from midphase
collision_worldid_out: collision world id from midphase
ncollision_out: number of collisions from broadphase and midphase
"""
pairid = wp.tid()
# only process pairs that are marked for heightfield collision (-1)
# the buffer is cleared at the start of each frame in collision_driver.py
# only process pairs that are marked for height field collision (-1)
if collision_hftri_index_in[pairid] != -1:
return
# get the collision pair info
pair = collision_pair_in[pairid]
# collision pair info
worldid = collision_worldid_in[pairid]
pair_id = collision_pairid_in[pairid]
# identify which geom is the heightfield
pair = collision_pair_in[pairid]
g1 = pair[0]
g2 = pair[1]
hfieldid = g1
geomid = g2
# if the first geom is not a heightfield, swap them
# in theory, shouldn't happen as _add_geom_pair already sorted the pair
# SHOULD NOT OCCUR: if the first geom is not a heightfield, swap
if geom_type[g1] != int(GeomType.HFIELD.value):
hfieldid = g2
geomid = g1
# get min/max grid coordinates for overlap region
min_i, min_j, max_i, max_j = _hfield_overlap_range(
geom_dataid,
geom_rbound,
geom_margin,
hfield_nrow,
hfield_ncol,
hfield_size,
geom_xpos_in,
geom_xmat_in,
hfieldid,
geomid,
worldid,
)
# height field info
hfdataid = geom_dataid[hfieldid]
size1 = hfield_size[hfdataid]
pos1 = geom_xpos_in[worldid, hfieldid]
mat1 = geom_xmat_in[worldid, hfieldid]
mat1T = wp.transpose(mat1)
# get hfield dimensions for triangle index calculation
dataid = geom_dataid[hfieldid]
ncol = hfield_ncol[dataid]
# geom info
pos2 = geom_xpos_in[worldid, geomid]
pos = mat1T @ (pos2 - pos1)
r2 = geom_rbound[worldid, geomid]
# loop through grid cells and add pairs for all triangles
for j in range(min_j, max_j + 1):
for i in range(min_i, max_i + 1):
# each grid cell contains two triangles
base_idx = ((j * (ncol - 1)) + i) * 2
# TODO(team): margin?
margin = wp.max(geom_margin[worldid, hfieldid], geom_margin[worldid, geomid])
# 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
# box-sphere test: vertical direction
if size1[2] < pos[2] - r2 - margin: # up
return
if -size1[3] > pos[2] + r2 + margin: # down
return
mat2 = geom_xmat_in[worldid, geomid]
mat = mat1T @ mat2
# aabb for geom in height field frame
xmax = -MJ_MAXVAL
ymax = -MJ_MAXVAL
zmax = -MJ_MAXVAL
xmin = MJ_MAXVAL
ymin = MJ_MAXVAL
zmin = MJ_MAXVAL
center2 = geom_aabb[geomid, 0]
size2 = geom_aabb[geomid, 1]
pos += mat1T @ center2
sign = wp.vec2(-1.0, 1.0)
for i in range(2):
for j in range(2):
for k in range(2):
corner_local = wp.vec3(sign[i] * size2[0], sign[j] * size2[1], sign[k] * size2[2])
corner_hf = mat @ corner_local
if corner_hf[0] > xmax:
xmax = corner_hf[0]
if corner_hf[1] > ymax:
ymax = corner_hf[1]
if corner_hf[2] > zmax:
zmax = corner_hf[2]
if corner_hf[0] < xmin:
xmin = corner_hf[0]
if corner_hf[1] < ymin:
ymin = corner_hf[1]
if corner_hf[2] < zmin:
zmin = corner_hf[2]
xmax += pos[0]
xmin += pos[0]
ymax += pos[1]
ymin += pos[1]
zmax += pos[2]
zmin += pos[2]
# box-box test
if (
(xmin - margin > size1[0])
or (xmax + margin < -size1[0])
or (ymin - margin > size1[1])
or (ymax + margin < -size1[1])
or (zmin - margin > size1[2])
or (zmax + margin < -size1[3])
):
return
# height field subgrid
nrow = hfield_nrow[hfieldid]
ncol = hfield_ncol[hfieldid]
size = hfield_size[hfieldid]
cmin, rmin, cmax, rmax = _hfield_subgrid(nrow, ncol, size, xmax, xmin, ymax, ymin)
# loop over subgrid triangles
for r in range(rmin, rmax):
for c in range(cmin, cmax):
# add both triangles from this cell
for t in range(2):
if i == 0 and j == 0 and t == 0:
for i in range(2):
if r == rmin and c == cmin and i == 0:
# reuse the initial pair for the 1st triangle
new_pairid = pairid
else:
# for the rest create a new pair
# create a new pair
new_pairid = wp.atomic_add(ncollision_out, 0, 1)
if new_pairid >= nconmax_in:
return
collision_pair_out[new_pairid] = pair
collision_hftri_index_out[new_pairid] = base_idx + t
collision_hftri_index_out[new_pairid] = 2 * (r * (ncol - 1) + c) + i
collision_pairid_out[new_pairid] = pair_id
collision_worldid_out[new_pairid] = worldid
def hfield_midphase(m: Model, d: Data):
"""Midphase collision detection for heightfield triangles with other geoms.
"""Midphase collision detection for height field triangles with other geoms.
Processes collision pairs from the broadphase where one geom is a heightfield and expands
Processes collision pairs from the broadphase where one geom is a height field and expands
them into multiple collision pairs, one for each potentially colliding triangle. The
function directly writes to the same collision buffers used by _add_geom_pair.
Args:
m: Model containing geometry and heightfield data
- geom_type: Array of geometry types
- geom_dataid: Array of geometry data IDs
- hfield_nrow: Array of heightfield rows
- hfield_ncol: Array of heightfield columns
- hfield_size: Array of heightfield sizes
- geom_rbound: Array of geometry bounding radii
- geom_margin: Array of geometry margins
d: Data containing current state and collision information
- nconmax: Maximum number of contacts
- geom_xpos: Array of geometry positions
- geom_xmat: Array of geometry orientation matrices
- collision_pair: Array of collision pairs
- collision_hftri_index: Array of heightfield triangle indices
- collision_pairid: Array of collision pair IDs
- collision_worldid: Array of collision world IDs
- ncollision: Number of collisions
"""
# launch the midphase kernel to expand height field collision pairs
# write directly to the same buffers that _add_geom_pair writes to
wp.launch(
kernel=_hfield_midphase,
dim=d.nconmax, # launch threads to process all potential pairs
dim=d.nconmax,
inputs=[
m.geom_type,
m.geom_dataid,
m.geom_aabb,
m.geom_rbound,
m.geom_margin,
m.hfield_nrow,
@@ -387,11 +397,5 @@ def hfield_midphase(m: Model, d: Data):
d.collision_pairid,
d.collision_worldid,
],
outputs=[
d.collision_pair,
d.collision_hftri_index,
d.collision_pairid,
d.collision_worldid,
d.ncollision,
],
outputs=[d.collision_pair, d.collision_hftri_index, d.collision_pairid, d.collision_worldid, d.ncollision],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+11 -25
View File
@@ -919,7 +919,7 @@ def plane_convex(
a_dist = wp.float32(-_HUGE_VAL)
for i in range(convex.vertnum):
support = wp.dot(plane_pos - convex.vert[convex.vertadr + i], n)
dist = wp.where(support > threshold, 0.0, -_HUGE_VAL)
dist = wp.where(support > threshold, support, -_HUGE_VAL)
if dist > a_dist:
indices[0] = i
a_dist = dist
@@ -942,7 +942,8 @@ def plane_convex(
for i in range(convex.vertnum):
support = wp.dot(plane_pos - convex.vert[convex.vertadr + i], n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
dist = wp.length_sq(ab - convex.vert[convex.vertadr + i]) + dist_mask
ap = a - convex.vert[convex.vertadr + i]
dist = wp.abs(wp.dot(ap, ab)) + dist_mask
if dist > c_dist:
indices[2] = i
c_dist = dist
@@ -955,8 +956,8 @@ def plane_convex(
for i in range(convex.vertnum):
support = wp.dot(plane_pos - convex.vert[convex.vertadr + i], n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
ap = ac - convex.vert[convex.vertadr + i]
bp = bc - convex.vert[convex.vertadr + i]
ap = a - convex.vert[convex.vertadr + i]
bp = b - convex.vert[convex.vertadr + i]
dist_ap = wp.abs(wp.dot(ap, ac)) + dist_mask
dist_bp = wp.abs(wp.dot(bp, bc)) + dist_mask
if dist_ap + dist_bp > d_dist:
@@ -993,10 +994,6 @@ def plane_convex(
threshold = wp.max(0.0, max_support - 1e-3)
a_dist = wp.float32(-_HUGE_VAL)
# hillclimb until no change
prev = int(-1)
imax = int(0)
while True:
prev = int(imax)
i = int(convex.graph[vert_edgeadr + imax])
@@ -1004,7 +1001,7 @@ def plane_convex(
subidx = convex.graph[edge_localid + i]
idx = convex.graph[vert_globalid + subidx]
support = wp.dot(plane_pos - convex.vert[convex.vertadr + idx], n)
dist = wp.where(support > threshold, 0.0, -_HUGE_VAL)
dist = wp.where(support > threshold, support, -_HUGE_VAL)
if dist > a_dist:
a_dist = dist
imax = int(subidx)
@@ -1017,10 +1014,6 @@ def plane_convex(
# Find point b (furthest from a)
b_dist = wp.float32(-_HUGE_VAL)
# hillclimb until no change
prev = int(-1)
imax = int(0)
while True:
prev = int(imax)
i = int(convex.graph[vert_edgeadr + imax])
@@ -1043,10 +1036,6 @@ def plane_convex(
# Find point c (furthest along axis orthogonal to a-b)
ab = wp.cross(n, a - b)
c_dist = wp.float32(-_HUGE_VAL)
# hillclimb until no change
prev = int(-1)
imax = int(0)
while True:
prev = int(imax)
i = int(convex.graph[vert_edgeadr + imax])
@@ -1055,7 +1044,8 @@ def plane_convex(
idx = convex.graph[vert_globalid + subidx]
support = wp.dot(plane_pos - convex.vert[convex.vertadr + idx], n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
dist = wp.length_sq(ab - convex.vert[convex.vertadr + idx]) + dist_mask
ap = a - convex.vert[convex.vertadr + i]
dist = wp.abs(wp.dot(ap, ab)) + dist_mask
if dist > c_dist:
c_dist = dist
imax = int(subidx)
@@ -1070,10 +1060,6 @@ def plane_convex(
ac = wp.cross(n, a - c)
bc = wp.cross(n, b - c)
d_dist = wp.float32(-_HUGE_VAL)
# hillclimb until no change
prev = int(-1)
imax = int(0)
while True:
prev = int(imax)
i = int(convex.graph[vert_edgeadr + imax])
@@ -1082,8 +1068,8 @@ def plane_convex(
idx = convex.graph[vert_globalid + subidx]
support = wp.dot(plane_pos - convex.vert[convex.vertadr + idx], n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
ap = ac - convex.vert[convex.vertadr + idx]
bp = bc - convex.vert[convex.vertadr + idx]
ap = a - convex.vert[convex.vertadr + idx]
bp = b - convex.vert[convex.vertadr + idx]
dist_ap = wp.abs(wp.dot(ap, ac)) + dist_mask
dist_bp = wp.abs(wp.dot(bp, bc)) + dist_mask
if dist_ap + dist_bp > d_dist:
@@ -2549,7 +2535,7 @@ def box_box(
v = (y * ax - x * ay) * C
if nl == 0:
if (u < 0 or u > 0) and (v < 0 or v > 1):
if (u < 0 or u > 1) and (v < 0 or v > 1):
continue
elif u < 0 or v < 0 or u > 1 or v > 1:
continue
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+3 -1
View File
@@ -312,6 +312,8 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
],
)
wp.copy(d.qacc_warmstart, d.qacc)
@wp.kernel
def _euler_damp_qfrc_sparse(
@@ -983,8 +985,8 @@ def forward(m: Model, d: Data):
energy = m.opt.enableflags & EnableBit.ENERGY
fwd_position(m, d, factorize=False)
d.sensordata.zero_()
sensor.sensor_pos(m, d)
if energy:
if m.sensor_e_potential == 0: # not computed by sensor
sensor.energy_pos(m, d)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py
Vendored
+1
View File
@@ -421,6 +421,7 @@ class ForwardTest(parameterized.TestCase):
"qfrc_actuator",
"qfrc_smooth",
"qacc",
"qacc_warmstart",
"qvel",
"qpos",
"efc_force",
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+5 -31
View File
@@ -396,13 +396,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
# contact sensor
sensor_adr_to_contact_adr = np.clip(np.cumsum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT) - 1, a_min=0, a_max=None)
# TODO(team): improve heuristic for selecting broadphase routine
if mjm.ngeom > 1000:
broadphase = types.BroadphaseType.SAP_SEGMENTED
elif mjm.ngeom > 100:
if nxn_geom_pair_filtered.shape[0] < 250_000:
broadphase = types.BroadphaseType.NXN
elif mjm.ngeom < 1000:
broadphase = types.BroadphaseType.SAP_TILE
else:
broadphase = types.BroadphaseType.NXN
broadphase = types.BroadphaseType.SAP_SEGMENTED
condim = np.concatenate((mjm.geom_condim, mjm.pair_dim))
condim_max = np.max(condim) if len(condim) > 0 else 0
@@ -473,6 +472,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
sdf_initpoints=mjm.opt.sdf_initpoints,
sdf_iterations=mjm.opt.sdf_iterations,
run_collision_detection=True,
legacy_gjk=False,
),
stat=types.Statistic(
meaninertia=mjm.stat.meaninertia,
@@ -1016,7 +1016,6 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
cost=wp.zeros((nworld,), dtype=float),
prev_cost=wp.zeros((nworld,), dtype=float),
state=wp.zeros((nworld, njmax), dtype=int),
gtol=wp.zeros((nworld,), dtype=float),
mv=wp.zeros((nworld, mjm.nv), dtype=float),
jv=wp.zeros((nworld, njmax), dtype=float),
quad=wp.zeros((nworld, njmax), dtype=wp.vec3f),
@@ -1028,18 +1027,6 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
beta=wp.zeros((nworld,), dtype=float),
done=wp.zeros((nworld,), dtype=bool),
# linesearch
ls_done=wp.zeros((nworld,), dtype=bool),
p0=wp.zeros((nworld,), dtype=wp.vec3),
lo=wp.zeros((nworld,), dtype=wp.vec3),
lo_alpha=wp.zeros((nworld,), dtype=float),
hi=wp.zeros((nworld,), dtype=wp.vec3),
hi_alpha=wp.zeros((nworld,), dtype=float),
lo_next=wp.zeros((nworld,), dtype=wp.vec3),
lo_next_alpha=wp.zeros((nworld,), dtype=float),
hi_next=wp.zeros((nworld,), dtype=wp.vec3),
hi_next_alpha=wp.zeros((nworld,), dtype=float),
mid=wp.zeros((nworld,), dtype=wp.vec3),
mid_alpha=wp.zeros((nworld,), dtype=float),
cost_candidate=wp.zeros((nworld, mjm.opt.ls_iterations), dtype=float),
),
# RK4
@@ -1395,7 +1382,6 @@ def put_data(
cost=wp.empty(shape=(nworld,), dtype=float),
prev_cost=wp.empty(shape=(nworld,), dtype=float),
state=wp.empty(shape=(nworld, njmax), dtype=int),
gtol=wp.empty(shape=(nworld,), dtype=float),
mv=wp.empty(shape=(nworld, mjm.nv), dtype=float),
jv=wp.empty(shape=(nworld, njmax), dtype=float),
quad=wp.empty(shape=(nworld, njmax), dtype=wp.vec3f),
@@ -1406,18 +1392,6 @@ def put_data(
prev_Mgrad=wp.empty(shape=(nworld, mjm.nv), dtype=float),
beta=wp.empty(shape=(nworld,), dtype=float),
done=wp.empty(shape=(nworld,), dtype=bool),
ls_done=wp.zeros(shape=(nworld,), dtype=bool),
p0=wp.empty(shape=(nworld,), dtype=wp.vec3),
lo=wp.empty(shape=(nworld,), dtype=wp.vec3),
lo_alpha=wp.empty(shape=(nworld,), dtype=float),
hi=wp.empty(shape=(nworld,), dtype=wp.vec3),
hi_alpha=wp.empty(shape=(nworld,), dtype=float),
lo_next=wp.empty(shape=(nworld,), dtype=wp.vec3),
lo_next_alpha=wp.empty(shape=(nworld,), dtype=float),
hi_next=wp.empty(shape=(nworld,), dtype=wp.vec3),
hi_next_alpha=wp.empty(shape=(nworld,), dtype=float),
mid=wp.empty(shape=(nworld,), dtype=wp.vec3),
mid_alpha=wp.empty(shape=(nworld,), dtype=float),
cost_candidate=wp.empty(shape=(nworld, mjm.opt.ls_iterations), dtype=float),
),
# TODO(team): skip allocation if integrator != RK4
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io_test.py
Vendored
+1 -1
View File
@@ -33,7 +33,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.io import MAX_WORLDS
_IO_TEST_MODELS = (
"pendula.xml",
"collision_sdf/tactile.xml",
"flex/cloth.xml",
"flex/floppy.xml",
"actuation/tendon_force_limit.xml",
"hfield/hfield.xml",
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/math.py
Vendored
+10
View File
@@ -190,6 +190,16 @@ def orthonormal(normal: wp.vec3) -> wp.vec3:
return dir
@wp.func
def orthonormal_to_z(normal: wp.vec3) -> wp.vec3:
if wp.abs(normal[0]) < wp.abs(normal[1]):
dir = wp.vec3(1.0 - normal[0] * normal[0], -normal[0] * normal[1], -normal[0] * normal[2])
else:
dir = wp.vec3(-normal[1] * normal[0], 1.0 - normal[1] * normal[1], -normal[1] * normal[2])
dir, _ = gjk_normalize(dir)
return dir
@wp.func
def gjk_normalize(a: wp.vec3):
norm = wp.length(a)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
-140
View File
@@ -226,18 +226,6 @@ def _ball_quat(jnt_qposadr: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float
return wp.normalize(quat)
@wp.kernel
def _limit_pos_zero(
# Model:
sensor_adr: wp.array(dtype=int),
sensor_limitpos_adr: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, limitposid = wp.tid()
sensordata_out[worldid, sensor_adr[sensor_limitpos_adr[limitposid]]] = 0.0
@wp.kernel
def _limit_pos(
# Model:
@@ -694,13 +682,6 @@ def sensor_pos(m: Model, d: Data):
)
# jointlimitpos and tendonlimitpos
wp.launch(
_limit_pos_zero,
dim=(d.nworld, m.sensor_limitpos_adr.size),
inputs=[m.sensor_adr, m.sensor_limitpos_adr],
outputs=[d.sensordata],
)
wp.launch(
_limit_pos,
dim=(d.nworld, d.njmax, m.sensor_limitpos_adr.size),
@@ -788,18 +769,6 @@ def _ball_ang_vel(jnt_dofadr: wp.array(dtype=int), qvel_in: wp.array2d(dtype=flo
return wp.vec3(qvel_in[worldid, adr + 0], qvel_in[worldid, adr + 1], qvel_in[worldid, adr + 2])
@wp.kernel
def _limit_vel_zero(
# Model:
sensor_adr: wp.array(dtype=int),
sensor_limitvel_adr: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, limitvelid = wp.tid()
sensordata_out[worldid, sensor_adr[sensor_limitvel_adr[limitvelid]]] = 0.0
@wp.kernel
def _limit_vel(
# Model:
@@ -1250,13 +1219,6 @@ def sensor_vel(m: Model, d: Data):
outputs=[d.sensordata],
)
wp.launch(
_limit_vel_zero,
dim=(d.nworld, m.sensor_limitvel_adr.size),
inputs=[m.sensor_adr, m.sensor_limitvel_adr],
outputs=[d.sensordata],
)
wp.launch(
_limit_vel,
dim=(d.nworld, d.njmax, m.sensor_limitvel_adr.size),
@@ -1367,20 +1329,6 @@ def _joint_actuator_force(
return qfrc_actuator_in[worldid, jnt_dofadr[objid]]
@wp.kernel
def _tendon_actuator_force_zero(
# Model:
sensor_adr: wp.array(dtype=int),
sensor_tendonactfrc_adr: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, tenactfrcid = wp.tid()
sensorid = sensor_tendonactfrc_adr[tenactfrcid]
adr = sensor_adr[sensorid]
sensordata_out[worldid, adr] = 0.0
@wp.kernel
def _tendon_actuator_force(
# Model:
@@ -1422,18 +1370,6 @@ def _tendon_actuator_force_cutoff(
_write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, sensordata_out[worldid])
@wp.kernel
def _limit_frc_zero(
# Model:
sensor_adr: wp.array(dtype=int),
sensor_limitfrc_adr: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, limitfrcid = wp.tid()
sensordata_out[worldid, sensor_adr[sensor_limitfrc_adr[limitfrcid]]] = 0.0
@wp.kernel
def _limit_frc(
# Model:
@@ -1760,20 +1696,6 @@ def _sensor_acc(
_write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
@wp.kernel
def _sensor_touch_zero(
# Model:
sensor_adr: wp.array(dtype=int),
sensor_touch_adr: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, sensortouchadrid = wp.tid()
sensorid = sensor_touch_adr[sensortouchadrid]
adr = sensor_adr[sensorid]
sensordata_out[worldid, adr] = 0.0
@wp.kernel
def _sensor_touch(
# Model:
@@ -1855,24 +1777,6 @@ def _sensor_touch(
wp.atomic_add(sensordata_out[worldid], adr, normalforce)
@wp.kernel
def _sensor_tactile_zero(
# Model:
sensor_type: wp.array(dtype=int),
sensor_dim: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, sensorid = wp.tid()
if sensor_type[sensorid] != int(SensorType.TACTILE.value):
return
for i in range(sensor_dim[sensorid]):
sensordata_out[worldid, sensor_adr[sensorid] + i] = 0.0
@wp.func
def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3:
return wp.spatial_bottom(vec) - wp.cross(dif, wp.spatial_top(vec))
@@ -2093,18 +1997,6 @@ def sensor_acc(m: Model, d: Data):
if m.opt.disableflags & DisableBit.SENSOR:
return
wp.launch(
_sensor_touch_zero,
dim=(d.nworld, m.sensor_touch_adr.size),
inputs=[
m.sensor_adr,
m.sensor_touch_adr,
],
outputs=[
d.sensordata,
],
)
wp.launch(
_sensor_touch,
dim=(d.nconmax, m.sensor_touch_adr.size),
@@ -2133,19 +2025,6 @@ def sensor_acc(m: Model, d: Data):
],
)
wp.launch(
_sensor_tactile_zero,
dim=(d.nworld, m.nsensor),
inputs=[
m.sensor_type,
m.sensor_dim,
m.sensor_adr,
],
outputs=[
d.sensordata,
],
)
wp.launch(
_sensor_tactile,
dim=(d.nconmax, m.nsensortaxel),
@@ -2280,18 +2159,6 @@ def sensor_acc(m: Model, d: Data):
outputs=[d.sensordata],
)
wp.launch(
_tendon_actuator_force_zero,
dim=(d.nworld, m.sensor_tendonactfrc_adr.size),
inputs=[
m.sensor_adr,
m.sensor_tendonactfrc_adr,
],
outputs=[
d.sensordata,
],
)
wp.launch(
_tendon_actuator_force,
dim=(d.nworld, m.sensor_tendonactfrc_adr.size, m.nu),
@@ -2321,13 +2188,6 @@ def sensor_acc(m: Model, d: Data):
outputs=[d.sensordata],
)
wp.launch(
_limit_frc_zero,
dim=(d.nworld, m.sensor_limitfrc_adr.size),
inputs=[m.sensor_adr, m.sensor_limitfrc_adr],
outputs=[d.sensordata],
)
wp.launch(
_limit_frc,
dim=(d.nworld, d.njmax, m.sensor_limitfrc_adr.size),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+57 -64
View File
@@ -347,10 +347,10 @@ def _subtree_com_init(
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
# Data out:
xipos_out: wp.array2d(dtype=wp.vec3),
subtree_com_out: wp.array2d(dtype=wp.vec3),
):
worldid, bodyid = wp.tid()
xipos_out[worldid, bodyid] = xipos_in[worldid, bodyid] * body_mass[worldid, bodyid]
subtree_com_out[worldid, bodyid] = xipos_in[worldid, bodyid] * body_mass[worldid, bodyid]
@wp.kernel
@@ -486,7 +486,7 @@ def com_pos(m: Model, d: Data):
Accumulates the mass-weighted positions up the kinematic tree, divides by total mass, and
computes composite inertias and motion degrees of freedom in the subtree CoM frame.
"""
wp.launch(_subtree_com_init, dim=(d.nworld, m.nbody), inputs=[m.body_mass, d.xipos, d.subtree_com])
wp.launch(_subtree_com_init, dim=(d.nworld, m.nbody), inputs=[m.body_mass, d.xipos], outputs=[d.subtree_com])
for i in reversed(range(len(m.body_tree))):
body_tree = m.body_tree[i]
@@ -514,37 +514,18 @@ def com_pos(m: Model, d: Data):
@wp.kernel
def _cam_local_to_global(
# Model:
cam_bodyid: wp.array(dtype=int),
cam_pos: wp.array2d(dtype=wp.vec3),
cam_quat: wp.array2d(dtype=wp.quat),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xquat_in: wp.array2d(dtype=wp.quat),
# Data out:
cam_xpos_out: wp.array2d(dtype=wp.vec3),
cam_xmat_out: wp.array2d(dtype=wp.mat33),
):
"""Fixed cameras."""
worldid, camid = wp.tid()
bodyid = cam_bodyid[camid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid, camid], xquat)
cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[worldid, camid]))
@wp.kernel
def _cam_fn(
# Model:
cam_mode: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
cam_targetbodyid: wp.array(dtype=int),
cam_pos: wp.array2d(dtype=wp.vec3),
cam_quat: wp.array2d(dtype=wp.quat),
cam_poscom0: wp.array2d(dtype=wp.vec3),
cam_pos0: wp.array2d(dtype=wp.vec3),
cam_mat0: wp.array2d(dtype=wp.mat33),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xquat_in: wp.array2d(dtype=wp.quat),
subtree_com_in: wp.array2d(dtype=wp.vec3),
# Data out:
cam_xpos_out: wp.array2d(dtype=wp.vec3),
@@ -556,7 +537,11 @@ def _cam_fn(
)
invalid_target = is_target_cam and (cam_targetbodyid[camid] < 0)
if invalid_target:
return
bodyid = cam_bodyid[camid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid, camid], xquat)
cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[worldid, camid]))
elif cam_mode[camid] == wp.static(CamLightType.TRACK.value):
cam_xmat_out[worldid, camid] = cam_mat0[worldid, camid]
body_xpos = xpos_in[worldid, cam_bodyid[camid]]
@@ -567,6 +552,10 @@ def _cam_fn(
elif cam_mode[camid] == wp.static(CamLightType.TARGETBODY.value) or cam_mode[camid] == wp.static(
CamLightType.TARGETBODYCOM.value
):
bodyid = cam_bodyid[camid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid, camid], xquat)
pos = xpos_in[worldid, cam_targetbodyid[camid]]
if cam_mode[camid] == wp.static(CamLightType.TARGETBODYCOM.value):
pos = subtree_com_in[worldid, cam_targetbodyid[camid]]
@@ -582,42 +571,28 @@ def _cam_fn(
mat_1[2], mat_2[2], mat_3[2]
)
# fmt: on
else:
bodyid = cam_bodyid[camid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid, camid], xquat)
cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[worldid, camid]))
@wp.kernel
def _light_local_to_global(
# Model:
light_bodyid: wp.array(dtype=int),
light_pos: wp.array2d(dtype=wp.vec3),
light_dir: wp.array2d(dtype=wp.vec3),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xquat_in: wp.array2d(dtype=wp.quat),
# Data out:
light_xpos_out: wp.array2d(dtype=wp.vec3),
light_xdir_out: wp.array2d(dtype=wp.vec3),
):
"""Fixed lights."""
worldid, lightid = wp.tid()
bodyid = light_bodyid[lightid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat)
light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[worldid, lightid], xquat)
@wp.kernel
def _light_fn(
# Model:
light_mode: wp.array(dtype=int),
light_bodyid: wp.array(dtype=int),
light_targetbodyid: wp.array(dtype=int),
light_pos: wp.array2d(dtype=wp.vec3),
light_dir: wp.array2d(dtype=wp.vec3),
light_poscom0: wp.array2d(dtype=wp.vec3),
light_pos0: wp.array2d(dtype=wp.vec3),
light_dir0: wp.array2d(dtype=wp.vec3),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
light_xpos_in: wp.array2d(dtype=wp.vec3),
xquat_in: wp.array2d(dtype=wp.quat),
subtree_com_in: wp.array2d(dtype=wp.vec3),
# Data out:
light_xpos_out: wp.array2d(dtype=wp.vec3),
@@ -629,6 +604,11 @@ def _light_fn(
)
invalid_target = is_target_light and (light_targetbodyid[lightid] < 0)
if invalid_target:
bodyid = light_bodyid[lightid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat)
light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[worldid, lightid], xquat)
return
elif light_mode[lightid] == wp.static(CamLightType.TRACK.value):
light_xdir_out[worldid, lightid] = light_dir0[worldid, lightid]
@@ -640,10 +620,21 @@ def _light_fn(
elif light_mode[lightid] == wp.static(CamLightType.TARGETBODY.value) or light_mode[lightid] == wp.static(
CamLightType.TARGETBODYCOM.value
):
bodyid = light_bodyid[lightid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat)
pos = xpos_in[worldid, light_targetbodyid[lightid]]
if light_mode[lightid] == wp.static(CamLightType.TARGETBODYCOM.value):
pos = subtree_com_in[worldid, light_targetbodyid[lightid]]
light_xdir_out[worldid, lightid] = pos - light_xpos_in[worldid, lightid]
light_xdir_out[worldid, lightid] = pos - light_xpos_out[worldid, lightid]
else:
bodyid = light_bodyid[lightid]
xpos = xpos_in[worldid, bodyid]
xquat = xquat_in[worldid, bodyid]
light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat)
light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[worldid, lightid], xquat)
light_xdir_out[worldid, lightid] = wp.normalize(light_xdir_out[worldid, lightid])
@@ -658,33 +649,35 @@ def camlight(m: Model, d: Data):
wp.launch(
_cam_local_to_global,
dim=(d.nworld, m.ncam),
inputs=[m.cam_bodyid, m.cam_pos, m.cam_quat, d.xpos, d.xquat],
outputs=[d.cam_xpos, d.cam_xmat],
)
wp.launch(
_cam_fn,
dim=(d.nworld, m.ncam),
inputs=[m.cam_mode, m.cam_bodyid, m.cam_targetbodyid, m.cam_poscom0, m.cam_pos0, m.cam_mat0, d.xpos, d.subtree_com],
inputs=[
m.cam_mode,
m.cam_bodyid,
m.cam_targetbodyid,
m.cam_pos,
m.cam_quat,
m.cam_poscom0,
m.cam_pos0,
m.cam_mat0,
d.xpos,
d.xquat,
d.subtree_com,
],
outputs=[d.cam_xpos, d.cam_xmat],
)
wp.launch(
_light_local_to_global,
dim=(d.nworld, m.nlight),
inputs=[m.light_bodyid, m.light_pos, m.light_dir, d.xpos, d.xquat],
outputs=[d.light_xpos, d.light_xdir],
)
wp.launch(
_light_fn,
dim=(d.nworld, m.nlight),
inputs=[
m.light_mode,
m.light_bodyid,
m.light_targetbodyid,
m.light_pos,
m.light_dir,
m.light_poscom0,
m.light_pos0,
m.light_dir0,
d.xpos,
d.light_xpos,
d.xquat,
d.subtree_com,
],
outputs=[d.light_xpos, d.light_xdir],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+362 -624
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver_test.py
Vendored
-2
View File
@@ -254,9 +254,7 @@ class SolverTest(parameterized.TestCase):
ls_parallel=ls_parallel,
)
qacc_warmstart = mjd.qacc_warmstart.copy()
mujoco.mj_forward(mjm, mjd)
mjd.qacc_warmstart = qacc_warmstart
d.qacc.zero_()
d.qfrc_constraint.zero_()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+3 -27
View File
@@ -55,7 +55,7 @@ class BlockDim:
cholesky_solve: int = 256
cholesky_factorize_solve: int = 256
# solver
update_gradient_cholesky: int = 256
update_gradient_cholesky: int = 64
# support
mul_m_dense: int = 256
@@ -570,6 +570,7 @@ class Option:
run_collision_detection: if False, skips collision detection and allows user-populated
contacts during the physics step (as opposed to DisableBit.CONTACT which explicitly
zeros out the contacts at each step)
legacy_gjk: run legacy gjk algorithm
"""
timestep: wp.array(dtype=float)
@@ -600,6 +601,7 @@ class Option:
sdf_initpoints: int
sdf_iterations: int
run_collision_detection: bool # warp only
legacy_gjk: bool
@dataclasses.dataclass
@@ -639,7 +641,6 @@ class Constraint:
cost: constraint + Gauss cost (nworld,)
prev_cost: cost from previous iter (nworld,)
state: constraint state (nworld, njmax)
gtol: linesearch termination tolerance (nworld,)
mv: qM @ search (nworld, nv)
jv: efc_J @ search (nworld, njmax)
quad: quadratic cost coefficients (nworld, njmax, 3)
@@ -650,18 +651,6 @@ class Constraint:
prev_Mgrad: previous Mgrad (nworld, nv)
beta: polak-ribiere beta (nworld,)
done: solver done (nworld,)
ls_done: linesearch done (nworld,)
p0: initial point (nworld, 3)
lo: low point bounding the line search interval (nworld, 3)
lo_alpha: alpha for low point (nworld,)
hi: high point bounding the line search interval (nworld, 3)
hi_alpha: alpha for high point (nworld,)
lo_next: next low point (nworld, 3)
lo_next_alpha: alpha for next low point (nworld,)
hi_next: next high point (nworld, 3)
hi_next_alpha: alpha for next high point (nworld,)
mid: loss at mid_alpha (nworld, 3)
mid_alpha: midpoint between lo_alpha and hi_alpha (nworld,)
cost_candidate: costs associated with step sizes (nworld, nlsp)
"""
@@ -688,7 +677,6 @@ class Constraint:
cost: wp.array(dtype=float)
prev_cost: wp.array(dtype=float)
state: wp.array2d(dtype=int)
gtol: wp.array(dtype=float)
mv: wp.array2d(dtype=float)
jv: wp.array2d(dtype=float)
quad: wp.array2d(dtype=wp.vec3)
@@ -700,18 +688,6 @@ class Constraint:
beta: wp.array(dtype=float)
done: wp.array(dtype=bool)
# linesearch
ls_done: wp.array(dtype=bool)
p0: wp.array(dtype=wp.vec3)
lo: wp.array(dtype=wp.vec3)
lo_alpha: wp.array(dtype=float)
hi: wp.array(dtype=wp.vec3)
hi_alpha: wp.array(dtype=float)
lo_next: wp.array(dtype=wp.vec3)
lo_next_alpha: wp.array(dtype=float)
hi_next: wp.array(dtype=wp.vec3)
hi_next_alpha: wp.array(dtype=float)
mid: wp.array(dtype=wp.vec3)
mid_alpha: wp.array(dtype=float)
cost_candidate: wp.array2d(dtype=float)
mjx/mujoco/mjx/third_party/mujoco_warp/test_data/flex/cloth.xml
Vendored
-45
View File
@@ -1,45 +0,0 @@
<!-- Copyright 2021 DeepMind Technologies Limited
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.
-->
<mujoco model="Poncho">
<include file="mannequin.xml"/>
<option timestep="0.001" integrator="Euler" solver="CG" tolerance="1e-6" jacobian="sparse">
<flag energy="enable"/>
</option>
<visual>
<map force="0.1" zfar="30"/>
<rgba haze="0.15 0.25 0.35 1"/>
<quality shadowsize="4096"/>
<global offwidth="800" offheight="800"/>
</visual>
<default>
<geom solref="0.003 1"/>
</default>
<worldbody>
<light directional="false" diffuse=".2 .2 .2" specular="0 0 0" pos="0 0 5" dir="0 0 -1"/>
<flexcomp name="towel" type="grid" count="15 15 1" spacing="0.1 0.1 0.1"
radius="0.03" dim="2" rgba="1 0.5 0.5 1" pos="0 0 2" mass=".1">
<edge equality="false"/>
<elasticity young="3e2" poisson="0" thickness="1e-1" damping="1e-3" elastic2d="both"/>
<contact vertcollide="true" conaffinity="0" contype="0"/>
</flexcomp>
</worldbody>
</mujoco>
mjx/mujoco/mjx/third_party/mujoco_warp/test_data/flex/mannequin.xml
Vendored
-174
View File
@@ -1,174 +0,0 @@
<!-- Copyright 2021 DeepMind Technologies Limited
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.
-->
<mujoco model="Mannequin">
<!-- Static humanoid model with all joints removed -->
<option timestep="0.005"/>
<visual>
<map force="0.1" zfar="30"/>
<rgba haze="0.15 0.25 0.35 1"/>
<global offwidth="2560" offheight="1440" elevation="-20" azimuth="120"/>
</visual>
<statistic center="0 0 0.7"/>
<asset>
<texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="32" height="512"/>
<texture name="body" type="cube" builtin="flat" mark="cross" width="128" height="128" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1" random="0.01"/>
<material name="body" texture="body" texuniform="true" rgba="0.8 0.6 .4 1"/>
<texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
<material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
</asset>
<default>
<motor ctrlrange="-1 1" ctrllimited="true"/>
<default class="body">
<!-- geoms -->
<geom type="capsule" condim="1" friction=".7" solimp=".9 .99 .003" solref=".003 1" material="body"/>
<default class="thigh">
<geom size=".06"/>
</default>
<default class="shin">
<geom fromto="0 0 0 0 0 -.3" size=".049"/>
</default>
<default class="foot">
<geom size=".027"/>
<default class="foot1">
<geom fromto="-.07 -.01 0 .14 -.03 0"/>
</default>
<default class="foot2">
<geom fromto="-.07 .01 0 .14 .03 0"/>
</default>
</default>
<default class="arm_upper">
<geom size=".04"/>
</default>
<default class="arm_lower">
<geom size=".031"/>
</default>
<default class="hand">
<geom type="sphere" size=".04"/>
</default>
<!-- joints -->
<joint type="hinge" damping=".2" stiffness="1" armature=".01" limited="true" solimplimit="0 .99 .01"/>
<default class="joint_big">
<joint damping="5" stiffness="10"/>
<default class="hip_x">
<joint range="-30 10"/>
</default>
<default class="hip_z">
<joint range="-60 35"/>
</default>
<default class="hip_y">
<joint axis="0 1 0" range="-150 20"/>
</default>
<default class="joint_big_stiff">
<joint stiffness="20"/>
</default>
</default>
<default class="knee">
<joint pos="0 0 .02" axis="0 -1 0" range="-160 2"/>
</default>
<default class="ankle">
<joint range="-50 50"/>
<default class="ankle_y">
<joint pos="0 0 .08" axis="0 1 0" stiffness="6"/>
</default>
<default class="ankle_x">
<joint pos="0 0 .04" stiffness="3"/>
</default>
</default>
<default class="shoulder">
<joint range="-85 60"/>
</default>
<default class="elbow">
<joint range="-100 50" stiffness="0"/>
</default>
</default>
</default>
<worldbody>
<geom name="floor" size="0 0 .05" type="plane" material="grid" condim="3"/>
<light name="spotlight" mode="targetbodycom" target="torso" diffuse=".8 .8 .8" specular="0.3 0.3 0.3" pos="0 -6 4" cutoff="30"/>
<body name="torso" pos="0 0 1.282" childclass="body">
<light name="top" pos="0 0 2" mode="trackcom"/>
<camera name="back" pos="-3 0 1" xyaxes="0 -1 0 1 0 2" mode="trackcom"/>
<camera name="side" pos="0 -3 1" xyaxes="1 0 0 0 1 2" mode="trackcom"/>
<freejoint name="root"/>
<geom name="torso" fromto="0 -.07 0 0 .07 0" size=".07"/>
<geom name="waist_upper" fromto="-.01 -.06 -.12 -.01 .06 -.12" size=".06"/>
<body name="head" pos="0 0 .19">
<geom name="head" type="sphere" size=".09"/>
<camera name="egocentric" pos=".09 0 0" xyaxes="0 -1 0 .1 0 1" fovy="80"/>
</body>
<body name="neck">
<geom type="cylinder" size=".03" fromto="0 0 .07 0 0 .1"/>
</body>
<body name="waist_lower" pos="-.01 0 -.26">
<geom name="waist_lower" fromto="0 -.06 0 0 .06 0" size=".06"/>
<body name="pelvis" pos="0 0 -.165">
<geom name="butt" fromto="-.02 -.07 0 -.02 .07 0" size=".09"/>
<body name="thigh_right" pos="0 -.1 -.04">
<geom name="thigh_right" fromto="0 0 0 0 .01 -.34" class="thigh"/>
<body name="shin_right" pos="0 .01 -.4">
<geom name="shin_right" class="shin"/>
<body name="foot_right" pos="0 0 -.39">
<geom name="foot1_right" class="foot1"/>
<geom name="foot2_right" class="foot2"/>
</body>
</body>
</body>
<body name="thigh_left" pos="0 .1 -.04">
<geom name="thigh_left" fromto="0 0 0 0 -.01 -.34" class="thigh"/>
<body name="shin_left" pos="0 -.01 -.4">
<geom name="shin_left" fromto="0 0 0 0 0 -.3" class="shin"/>
<body name="foot_left" pos="0 0 -.39">
<geom name="foot1_left" class="foot1"/>
<geom name="foot2_left" class="foot2"/>
</body>
</body>
</body>
</body>
</body>
<body name="upper_arm_right" pos="0 -.17 .06">
<geom name="upper_arm_right" fromto="0 0 0 0 -.22 0" class="arm_upper"/>
<body name="lower_arm_right" pos="0 -.22 0">
<geom name="lower_arm_right" fromto=".01 .01 .01 0 -.21 -.05" class="arm_lower"/>
<body name="hand_right" pos="0 -.22 -.05">
<geom name="hand_right" zaxis="1 1 1" class="hand"/>
</body>
</body>
</body>
<body name="upper_arm_left" pos="0 .17 .06">
<geom name="upper_arm_left" fromto="0 0 0 0 .22 0" class="arm_upper"/>
<body name="lower_arm_left" pos="0 .22 0">
<geom name="lower_arm_left" fromto=".01 -.01 .01 0 .21 -.05" class="arm_lower"/>
<body name="hand_left" pos="0 .22 -.05">
<geom name="hand_left" zaxis="1 -1 1" class="hand"/>
</body>
</body>
</body>
</body>
</worldbody>
<contact>
<exclude body1="waist_lower" body2="thigh_right"/>
<exclude body1="waist_lower" body2="thigh_left"/>
</contact>
</mujoco>
mjx/mujoco/mjx/warp/collision_driver.py
+4
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
@@ -105,6 +106,7 @@ def _collision_shim(
opt__epa_iterations: int,
opt__gjk_iterations: int,
opt__graph_conditional: bool,
opt__legacy_gjk: bool,
opt__sdf_initpoints: int,
opt__sdf_iterations: int,
# Data
@@ -201,6 +203,7 @@ def _collision_shim(
_m.opt.epa_iterations = opt__epa_iterations
_m.opt.gjk_iterations = opt__gjk_iterations
_m.opt.graph_conditional = opt__graph_conditional
_m.opt.legacy_gjk = opt__legacy_gjk
_m.opt.sdf_initpoints = opt__sdf_initpoints
_m.opt.sdf_iterations = opt__sdf_iterations
_m.pair_dim = pair_dim
@@ -400,6 +403,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.opt._impl.epa_iterations,
m.opt._impl.gjk_iterations,
m.opt._impl.graph_conditional,
m.opt._impl.legacy_gjk,
m.opt._impl.sdf_initpoints,
m.opt._impl.sdf_iterations,
d._impl.nconmax,
mjx/mujoco/mjx/warp/forward.py
+40 -196
View File
@@ -227,7 +227,6 @@ def _forward_shim(
nlsp: int,
nmeshface: int,
nmocap: int,
nsensor: int,
nsensortaxel: int,
nsite: int,
ntendon: int,
@@ -331,6 +330,7 @@ def _forward_shim(
opt__impratio: wp.array(dtype=float),
opt__is_sparse: bool,
opt__iterations: int,
opt__legacy_gjk: bool,
opt__ls_iterations: int,
opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
@@ -496,24 +496,11 @@ def _forward_shim(
efc__gauss: wp.array(dtype=float),
efc__grad: wp.array2d(dtype=float),
efc__grad_dot: wp.array(dtype=float),
efc__gtol: wp.array(dtype=float),
efc__h: wp.array3d(dtype=float),
efc__hi: wp.array(dtype=wp.vec3),
efc__hi_alpha: wp.array(dtype=float),
efc__hi_next: wp.array(dtype=wp.vec3),
efc__hi_next_alpha: wp.array(dtype=float),
efc__id: wp.array2d(dtype=int),
efc__jv: wp.array2d(dtype=float),
efc__lo: wp.array(dtype=wp.vec3),
efc__lo_alpha: wp.array(dtype=float),
efc__lo_next: wp.array(dtype=wp.vec3),
efc__lo_next_alpha: wp.array(dtype=float),
efc__ls_done: wp.array(dtype=bool),
efc__margin: wp.array2d(dtype=float),
efc__mid: wp.array(dtype=wp.vec3),
efc__mid_alpha: wp.array(dtype=float),
efc__mv: wp.array2d(dtype=float),
efc__p0: wp.array(dtype=wp.vec3),
efc__pos: wp.array2d(dtype=float),
efc__prev_Mgrad: wp.array2d(dtype=float),
efc__prev_cost: wp.array(dtype=float),
@@ -710,7 +697,6 @@ def _forward_shim(
_m.nlsp = nlsp
_m.nmeshface = nmeshface
_m.nmocap = nmocap
_m.nsensor = nsensor
_m.nsensortaxel = nsensortaxel
_m.nsite = nsite
_m.ntendon = ntendon
@@ -733,6 +719,7 @@ def _forward_shim(
_m.opt.impratio = opt__impratio
_m.opt.is_sparse = opt__is_sparse
_m.opt.iterations = opt__iterations
_m.opt.legacy_gjk = opt__legacy_gjk
_m.opt.ls_iterations = opt__ls_iterations
_m.opt.ls_parallel = opt__ls_parallel
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
@@ -880,24 +867,11 @@ def _forward_shim(
_d.efc.gauss = efc__gauss
_d.efc.grad = efc__grad
_d.efc.grad_dot = efc__grad_dot
_d.efc.gtol = efc__gtol
_d.efc.h = efc__h
_d.efc.hi = efc__hi
_d.efc.hi_alpha = efc__hi_alpha
_d.efc.hi_next = efc__hi_next
_d.efc.hi_next_alpha = efc__hi_next_alpha
_d.efc.id = efc__id
_d.efc.jv = efc__jv
_d.efc.lo = efc__lo
_d.efc.lo_alpha = efc__lo_alpha
_d.efc.lo_next = efc__lo_next
_d.efc.lo_next_alpha = efc__lo_next_alpha
_d.efc.ls_done = efc__ls_done
_d.efc.margin = efc__margin
_d.efc.mid = efc__mid
_d.efc.mid_alpha = efc__mid_alpha
_d.efc.mv = efc__mv
_d.efc.p0 = efc__p0
_d.efc.pos = efc__pos
_d.efc.prev_Mgrad = efc__prev_Mgrad
_d.efc.prev_cost = efc__prev_cost
@@ -1161,24 +1135,11 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'efc__gauss': d._impl.efc__gauss.shape,
'efc__grad': d._impl.efc__grad.shape,
'efc__grad_dot': d._impl.efc__grad_dot.shape,
'efc__gtol': d._impl.efc__gtol.shape,
'efc__h': d._impl.efc__h.shape,
'efc__hi': d._impl.efc__hi.shape,
'efc__hi_alpha': d._impl.efc__hi_alpha.shape,
'efc__hi_next': d._impl.efc__hi_next.shape,
'efc__hi_next_alpha': d._impl.efc__hi_next_alpha.shape,
'efc__id': d._impl.efc__id.shape,
'efc__jv': d._impl.efc__jv.shape,
'efc__lo': d._impl.efc__lo.shape,
'efc__lo_alpha': d._impl.efc__lo_alpha.shape,
'efc__lo_next': d._impl.efc__lo_next.shape,
'efc__lo_next_alpha': d._impl.efc__lo_next_alpha.shape,
'efc__ls_done': d._impl.efc__ls_done.shape,
'efc__margin': d._impl.efc__margin.shape,
'efc__mid': d._impl.efc__mid.shape,
'efc__mid_alpha': d._impl.efc__mid_alpha.shape,
'efc__mv': d._impl.efc__mv.shape,
'efc__p0': d._impl.efc__p0.shape,
'efc__pos': d._impl.efc__pos.shape,
'efc__prev_Mgrad': d._impl.efc__prev_Mgrad.shape,
'efc__prev_cost': d._impl.efc__prev_cost.shape,
@@ -1193,7 +1154,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_forward_shim,
num_outputs=177,
num_outputs=164,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
@@ -1345,24 +1306,11 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'efc__gauss',
'efc__grad',
'efc__grad_dot',
'efc__gtol',
'efc__h',
'efc__hi',
'efc__hi_alpha',
'efc__hi_next',
'efc__hi_next_alpha',
'efc__id',
'efc__jv',
'efc__lo',
'efc__lo_alpha',
'efc__lo_next',
'efc__lo_next_alpha',
'efc__ls_done',
'efc__margin',
'efc__mid',
'efc__mid_alpha',
'efc__mv',
'efc__p0',
'efc__pos',
'efc__prev_Mgrad',
'efc__prev_cost',
@@ -1558,7 +1506,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.nlsp,
m.nmeshface,
m.nmocap,
m.nsensor,
m._impl.nsensortaxel,
m.nsite,
m.ntendon,
@@ -1662,6 +1609,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.opt.impratio,
m.opt._impl.is_sparse,
m.opt.iterations,
m.opt._impl.legacy_gjk,
m.opt.ls_iterations,
m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
@@ -1826,24 +1774,11 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.efc__gauss,
d._impl.efc__grad,
d._impl.efc__grad_dot,
d._impl.efc__gtol,
d._impl.efc__h,
d._impl.efc__hi,
d._impl.efc__hi_alpha,
d._impl.efc__hi_next,
d._impl.efc__hi_next_alpha,
d._impl.efc__id,
d._impl.efc__jv,
d._impl.efc__lo,
d._impl.efc__lo_alpha,
d._impl.efc__lo_next,
d._impl.efc__lo_next_alpha,
d._impl.efc__ls_done,
d._impl.efc__margin,
d._impl.efc__mid,
d._impl.efc__mid_alpha,
d._impl.efc__mv,
d._impl.efc__p0,
d._impl.efc__pos,
d._impl.efc__prev_Mgrad,
d._impl.efc__prev_cost,
@@ -2005,35 +1940,22 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'_impl.efc__gauss': out[145],
'_impl.efc__grad': out[146],
'_impl.efc__grad_dot': out[147],
'_impl.efc__gtol': out[148],
'_impl.efc__h': out[149],
'_impl.efc__hi': out[150],
'_impl.efc__hi_alpha': out[151],
'_impl.efc__hi_next': out[152],
'_impl.efc__hi_next_alpha': out[153],
'_impl.efc__id': out[154],
'_impl.efc__jv': out[155],
'_impl.efc__lo': out[156],
'_impl.efc__lo_alpha': out[157],
'_impl.efc__lo_next': out[158],
'_impl.efc__lo_next_alpha': out[159],
'_impl.efc__ls_done': out[160],
'_impl.efc__margin': out[161],
'_impl.efc__mid': out[162],
'_impl.efc__mid_alpha': out[163],
'_impl.efc__mv': out[164],
'_impl.efc__p0': out[165],
'_impl.efc__pos': out[166],
'_impl.efc__prev_Mgrad': out[167],
'_impl.efc__prev_cost': out[168],
'_impl.efc__prev_grad': out[169],
'_impl.efc__quad': out[170],
'_impl.efc__quad_gauss': out[171],
'_impl.efc__search': out[172],
'_impl.efc__search_dot': out[173],
'_impl.efc__state': out[174],
'_impl.efc__type': out[175],
'_impl.efc__vel': out[176],
'_impl.efc__h': out[148],
'_impl.efc__id': out[149],
'_impl.efc__jv': out[150],
'_impl.efc__margin': out[151],
'_impl.efc__mv': out[152],
'_impl.efc__pos': out[153],
'_impl.efc__prev_Mgrad': out[154],
'_impl.efc__prev_cost': out[155],
'_impl.efc__prev_grad': out[156],
'_impl.efc__quad': out[157],
'_impl.efc__quad_gauss': out[158],
'_impl.efc__search': out[159],
'_impl.efc__search_dot': out[160],
'_impl.efc__state': out[161],
'_impl.efc__type': out[162],
'_impl.efc__vel': out[163],
})
return d
@@ -2254,7 +2176,6 @@ def _step_shim(
nlsp: int,
nmeshface: int,
nmocap: int,
nsensor: int,
nsensortaxel: int,
nsite: int,
ntendon: int,
@@ -2359,6 +2280,7 @@ def _step_shim(
opt__integrator: int,
opt__is_sparse: bool,
opt__iterations: int,
opt__legacy_gjk: bool,
opt__ls_iterations: int,
opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
@@ -2536,24 +2458,11 @@ def _step_shim(
efc__gauss: wp.array(dtype=float),
efc__grad: wp.array2d(dtype=float),
efc__grad_dot: wp.array(dtype=float),
efc__gtol: wp.array(dtype=float),
efc__h: wp.array3d(dtype=float),
efc__hi: wp.array(dtype=wp.vec3),
efc__hi_alpha: wp.array(dtype=float),
efc__hi_next: wp.array(dtype=wp.vec3),
efc__hi_next_alpha: wp.array(dtype=float),
efc__id: wp.array2d(dtype=int),
efc__jv: wp.array2d(dtype=float),
efc__lo: wp.array(dtype=wp.vec3),
efc__lo_alpha: wp.array(dtype=float),
efc__lo_next: wp.array(dtype=wp.vec3),
efc__lo_next_alpha: wp.array(dtype=float),
efc__ls_done: wp.array(dtype=bool),
efc__margin: wp.array2d(dtype=float),
efc__mid: wp.array(dtype=wp.vec3),
efc__mid_alpha: wp.array(dtype=float),
efc__mv: wp.array2d(dtype=float),
efc__p0: wp.array(dtype=wp.vec3),
efc__pos: wp.array2d(dtype=float),
efc__prev_Mgrad: wp.array2d(dtype=float),
efc__prev_cost: wp.array(dtype=float),
@@ -2751,7 +2660,6 @@ def _step_shim(
_m.nlsp = nlsp
_m.nmeshface = nmeshface
_m.nmocap = nmocap
_m.nsensor = nsensor
_m.nsensortaxel = nsensortaxel
_m.nsite = nsite
_m.ntendon = ntendon
@@ -2775,6 +2683,7 @@ def _step_shim(
_m.opt.integrator = opt__integrator
_m.opt.is_sparse = opt__is_sparse
_m.opt.iterations = opt__iterations
_m.opt.legacy_gjk = opt__legacy_gjk
_m.opt.ls_iterations = opt__ls_iterations
_m.opt.ls_parallel = opt__ls_parallel
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
@@ -2924,24 +2833,11 @@ def _step_shim(
_d.efc.gauss = efc__gauss
_d.efc.grad = efc__grad
_d.efc.grad_dot = efc__grad_dot
_d.efc.gtol = efc__gtol
_d.efc.h = efc__h
_d.efc.hi = efc__hi
_d.efc.hi_alpha = efc__hi_alpha
_d.efc.hi_next = efc__hi_next
_d.efc.hi_next_alpha = efc__hi_next_alpha
_d.efc.id = efc__id
_d.efc.jv = efc__jv
_d.efc.lo = efc__lo
_d.efc.lo_alpha = efc__lo_alpha
_d.efc.lo_next = efc__lo_next
_d.efc.lo_next_alpha = efc__lo_next_alpha
_d.efc.ls_done = efc__ls_done
_d.efc.margin = efc__margin
_d.efc.mid = efc__mid
_d.efc.mid_alpha = efc__mid_alpha
_d.efc.mv = efc__mv
_d.efc.p0 = efc__p0
_d.efc.pos = efc__pos
_d.efc.prev_Mgrad = efc__prev_Mgrad
_d.efc.prev_cost = efc__prev_cost
@@ -3227,24 +3123,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'efc__gauss': d._impl.efc__gauss.shape,
'efc__grad': d._impl.efc__grad.shape,
'efc__grad_dot': d._impl.efc__grad_dot.shape,
'efc__gtol': d._impl.efc__gtol.shape,
'efc__h': d._impl.efc__h.shape,
'efc__hi': d._impl.efc__hi.shape,
'efc__hi_alpha': d._impl.efc__hi_alpha.shape,
'efc__hi_next': d._impl.efc__hi_next.shape,
'efc__hi_next_alpha': d._impl.efc__hi_next_alpha.shape,
'efc__id': d._impl.efc__id.shape,
'efc__jv': d._impl.efc__jv.shape,
'efc__lo': d._impl.efc__lo.shape,
'efc__lo_alpha': d._impl.efc__lo_alpha.shape,
'efc__lo_next': d._impl.efc__lo_next.shape,
'efc__lo_next_alpha': d._impl.efc__lo_next_alpha.shape,
'efc__ls_done': d._impl.efc__ls_done.shape,
'efc__margin': d._impl.efc__margin.shape,
'efc__mid': d._impl.efc__mid.shape,
'efc__mid_alpha': d._impl.efc__mid_alpha.shape,
'efc__mv': d._impl.efc__mv.shape,
'efc__p0': d._impl.efc__p0.shape,
'efc__pos': d._impl.efc__pos.shape,
'efc__prev_Mgrad': d._impl.efc__prev_Mgrad.shape,
'efc__prev_cost': d._impl.efc__prev_cost.shape,
@@ -3259,7 +3142,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_step_shim,
num_outputs=189,
num_outputs=176,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
@@ -3423,24 +3306,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'efc__gauss',
'efc__grad',
'efc__grad_dot',
'efc__gtol',
'efc__h',
'efc__hi',
'efc__hi_alpha',
'efc__hi_next',
'efc__hi_next_alpha',
'efc__id',
'efc__jv',
'efc__lo',
'efc__lo_alpha',
'efc__lo_next',
'efc__lo_next_alpha',
'efc__ls_done',
'efc__margin',
'efc__mid',
'efc__mid_alpha',
'efc__mv',
'efc__p0',
'efc__pos',
'efc__prev_Mgrad',
'efc__prev_cost',
@@ -3637,7 +3507,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.nlsp,
m.nmeshface,
m.nmocap,
m.nsensor,
m._impl.nsensortaxel,
m.nsite,
m.ntendon,
@@ -3742,6 +3611,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.opt.integrator,
m.opt._impl.is_sparse,
m.opt.iterations,
m.opt._impl.legacy_gjk,
m.opt.ls_iterations,
m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
@@ -3918,24 +3788,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.efc__gauss,
d._impl.efc__grad,
d._impl.efc__grad_dot,
d._impl.efc__gtol,
d._impl.efc__h,
d._impl.efc__hi,
d._impl.efc__hi_alpha,
d._impl.efc__hi_next,
d._impl.efc__hi_next_alpha,
d._impl.efc__id,
d._impl.efc__jv,
d._impl.efc__lo,
d._impl.efc__lo_alpha,
d._impl.efc__lo_next,
d._impl.efc__lo_next_alpha,
d._impl.efc__ls_done,
d._impl.efc__margin,
d._impl.efc__mid,
d._impl.efc__mid_alpha,
d._impl.efc__mv,
d._impl.efc__p0,
d._impl.efc__pos,
d._impl.efc__prev_Mgrad,
d._impl.efc__prev_cost,
@@ -4109,35 +3966,22 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'_impl.efc__gauss': out[157],
'_impl.efc__grad': out[158],
'_impl.efc__grad_dot': out[159],
'_impl.efc__gtol': out[160],
'_impl.efc__h': out[161],
'_impl.efc__hi': out[162],
'_impl.efc__hi_alpha': out[163],
'_impl.efc__hi_next': out[164],
'_impl.efc__hi_next_alpha': out[165],
'_impl.efc__id': out[166],
'_impl.efc__jv': out[167],
'_impl.efc__lo': out[168],
'_impl.efc__lo_alpha': out[169],
'_impl.efc__lo_next': out[170],
'_impl.efc__lo_next_alpha': out[171],
'_impl.efc__ls_done': out[172],
'_impl.efc__margin': out[173],
'_impl.efc__mid': out[174],
'_impl.efc__mid_alpha': out[175],
'_impl.efc__mv': out[176],
'_impl.efc__p0': out[177],
'_impl.efc__pos': out[178],
'_impl.efc__prev_Mgrad': out[179],
'_impl.efc__prev_cost': out[180],
'_impl.efc__prev_grad': out[181],
'_impl.efc__quad': out[182],
'_impl.efc__quad_gauss': out[183],
'_impl.efc__search': out[184],
'_impl.efc__search_dot': out[185],
'_impl.efc__state': out[186],
'_impl.efc__type': out[187],
'_impl.efc__vel': out[188],
'_impl.efc__h': out[160],
'_impl.efc__id': out[161],
'_impl.efc__jv': out[162],
'_impl.efc__margin': out[163],
'_impl.efc__mv': out[164],
'_impl.efc__pos': out[165],
'_impl.efc__prev_Mgrad': out[166],
'_impl.efc__prev_cost': out[167],
'_impl.efc__prev_grad': out[168],
'_impl.efc__quad': out[169],
'_impl.efc__quad_gauss': out[170],
'_impl.efc__search': out[171],
'_impl.efc__search_dot': out[172],
'_impl.efc__state': out[173],
'_impl.efc__type': out[174],
'_impl.efc__vel': out[175],
})
return d
mjx/mujoco/mjx/warp/types.py
+5 -39
View File
@@ -94,6 +94,7 @@ class OptionWarp(PyTreeNode):
graph_conditional: bool
has_fluid: bool
is_sparse: bool
legacy_gjk: bool
ls_parallel: bool
ls_parallel_min_step: float
run_collision_detection: bool
@@ -256,24 +257,11 @@ class DataWarp(PyTreeNode):
efc__gauss: jax.Array
efc__grad: jax.Array
efc__grad_dot: jax.Array
efc__gtol: jax.Array
efc__h: jax.Array
efc__hi: jax.Array
efc__hi_alpha: jax.Array
efc__hi_next: jax.Array
efc__hi_next_alpha: jax.Array
efc__id: jax.Array
efc__jv: jax.Array
efc__lo: jax.Array
efc__lo_alpha: jax.Array
efc__lo_next: jax.Array
efc__lo_next_alpha: jax.Array
efc__ls_done: jax.Array
efc__margin: jax.Array
efc__mid: jax.Array
efc__mid_alpha: jax.Array
efc__mv: jax.Array
efc__p0: jax.Array
efc__pos: jax.Array
efc__prev_Mgrad: jax.Array
efc__prev_cost: jax.Array
@@ -488,24 +476,11 @@ _NDIM = {
'efc__gauss': 1,
'efc__grad': 2,
'efc__grad_dot': 1,
'efc__gtol': 1,
'efc__h': 3,
'efc__hi': 2,
'efc__hi_alpha': 1,
'efc__hi_next': 2,
'efc__hi_next_alpha': 1,
'efc__id': 2,
'efc__jv': 2,
'efc__lo': 2,
'efc__lo_alpha': 1,
'efc__lo_next': 2,
'efc__lo_next_alpha': 1,
'efc__ls_done': 1,
'efc__margin': 2,
'efc__mid': 2,
'efc__mid_alpha': 1,
'efc__mv': 2,
'efc__p0': 2,
'efc__pos': 2,
'efc__prev_Mgrad': 2,
'efc__prev_cost': 1,
@@ -884,6 +859,7 @@ _NDIM = {
'opt__integrator': 0,
'opt__is_sparse': 0,
'opt__iterations': 0,
'opt__legacy_gjk': 0,
'opt__ls_iterations': 0,
'opt__ls_parallel': 0,
'opt__ls_parallel_min_step': 0,
@@ -1002,6 +978,7 @@ _NDIM = {
'integrator': 0,
'is_sparse': 0,
'iterations': 0,
'legacy_gjk': 0,
'ls_iterations': 0,
'ls_parallel': 0,
'ls_parallel_min_step': 0,
@@ -1075,24 +1052,11 @@ _BATCH_DIM = {
'efc__gauss': True,
'efc__grad': True,
'efc__grad_dot': True,
'efc__gtol': True,
'efc__h': True,
'efc__hi': True,
'efc__hi_alpha': True,
'efc__hi_next': True,
'efc__hi_next_alpha': True,
'efc__id': True,
'efc__jv': True,
'efc__lo': True,
'efc__lo_alpha': True,
'efc__lo_next': True,
'efc__lo_next_alpha': True,
'efc__ls_done': True,
'efc__margin': True,
'efc__mid': True,
'efc__mid_alpha': True,
'efc__mv': True,
'efc__p0': True,
'efc__pos': True,
'efc__prev_Mgrad': True,
'efc__prev_cost': True,
@@ -1471,6 +1435,7 @@ _BATCH_DIM = {
'opt__integrator': False,
'opt__is_sparse': False,
'opt__iterations': False,
'opt__legacy_gjk': False,
'opt__ls_iterations': False,
'opt__ls_parallel': False,
'opt__ls_parallel_min_step': False,
@@ -1589,6 +1554,7 @@ _BATCH_DIM = {
'integrator': False,
'is_sparse': False,
'iterations': False,
'legacy_gjk': False,
'ls_iterations': False,
'ls_parallel': False,
'ls_parallel_min_step': False,