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

Import google-deepmind/mujoco_warp from GitHub. Update MJX io.py.

Browse Source 
PiperOrigin-RevId: 816431884
Change-Id: I54881e7e350fbe5b2be1833e66745894b122f374
This commit is contained in:
Baruch Tabanpour
2025-10-07 16:42:44 -07:00
committed by Copybara-Service
parent f83ee3fd73
commit b56cf98e41
25 changed files with 1026 additions and 549 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+3 -1
View File
@@ -30,11 +30,13 @@ General
:class: attention
- The attribute ``mjOption.apirate`` was unused and has been removed.
- MJX ``nconmax`` and ``njmax`` fields in ``mjx.make_data`` now default to ``None`` instead of -1.
MJX
^^^
- Fix :github:issue:`2508`, ``qLD`` shapes mismatched mjModel during ``get_data_into``.
- Pull in MuJoCo Warp updatest to ``io.py``, and use ``naconmax`` instead of ``naconmax`` to set the maximum number of
contacts over all environments.
Version 3.3.6 (September 15, 2025)
----------------------------------
mjx/mujoco/mjx/_src/io.py
+12 -8
View File
@@ -827,8 +827,8 @@ def _get_nested_attr(obj: Any, attr_name: str, split: str) -> Any:
def _make_data_warp(
m: Union[types.Model, mujoco.MjModel],
device: Optional[jax.Device] = None,
nconmax: int = -1,
njmax: int = -1,
nconmax: Optional[int] = None,
njmax: Optional[int] = None,
) -> types.Data:
"""Allocate and initialize Data for the Warp implementation."""
if not isinstance(m, mujoco.MjModel):
@@ -841,7 +841,7 @@ def _make_data_warp(
raise RuntimeError('Warp is not installed.')
with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable
dw = mjwp.make_data(m, nworld=1, nconmax=nconmax, njmax=njmax) # pylint: disable=undefined-variable
dw = mjwp.make_data(m, nworld=1, naconmax=nconmax, njmax=njmax) # pylint: disable=undefined-variable
fields = _make_data_public_fields(m)
for k in fields:
@@ -876,7 +876,7 @@ def _make_data_warp(
# TODO(robotics-simulation): remove this warmup compilation once warp
# stops unloading modules during XLA graph capture for tile kernels.
# pylint: disable=undefined-variable
dw = mjwp.make_data(m, nworld=1)
dw = mjwp.make_data(m, nworld=1, naconmax=nconmax, njmax=njmax)
mw = mjwp.put_model(m)
_ = mjwp.step(mw, dw)
# pylint: enable=undefined-variable
@@ -890,8 +890,8 @@ def make_data(
device: Optional[jax.Device] = None,
impl: Optional[Union[str, types.Impl]] = None,
_full_compat: bool = False, # pylint: disable=invalid-name
nconmax: int = -1,
njmax: int = -1,
nconmax: Optional[int] = None,
njmax: Optional[int] = None,
) -> types.Data:
"""Allocate and initialize Data.
@@ -902,8 +902,12 @@ def make_data(
_full_compat: put all fields onto device irrespective of MJX support This is
an experimental feature. Avoid using it for now. If using this flag, also
use _full_compat for put_model.
nconmax: maximum number of contacts to allocate for warp
njmax: maximum number of constraints to allocate for warp
nconmax: maximum number of contacts to allocate for warp across all worlds
Since the number of worlds is **not** pre-defined in JAX, we use the
`nconmax` argument to set the upper bound for the number of contacts
across all worlds. In MuJoCo Warp, the analgous field is called
`naconmax`.
njmax: maximum number of constraints to allocate for warp across all worlds
Returns:
an initialized mjx.Data placed on device
mjx/mujoco/mjx/_src/io_test.py
+2 -2
View File
@@ -786,8 +786,8 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(dx[0].qpos.shape, (m.nq,))
if impl == 'warp':
self.assertEqual(dx._impl.contact__dist.shape, (dx._impl.nconmax,))
self.assertEqual(dx[0]._impl.contact__dist.shape, (dx._impl.nconmax,))
self.assertEqual(dx._impl.contact__dist.shape, (dx._impl.naconmax,))
self.assertEqual(dx[0]._impl.contact__dist.shape, (dx._impl.naconmax,))
class FullCompatTest(parameterized.TestCase):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
Vendored
+3 -3
View File
@@ -105,7 +105,7 @@ def benchmark(
"""
trace = {}
ncon, nefc, solver_niter = [], [], []
nacon, nefc, solver_niter = [], [], []
center = wp.array([], dtype=wp.float32)
with warp_util.EventTracer(enabled=event_trace) as tracer:
@@ -142,7 +142,7 @@ def benchmark(
else:
trace = tracer.trace()
if measure_alloc:
ncon.append(np.max([d.ncon.numpy()[0], d.ncollision.numpy()[0]]))
nacon.append(np.max([d.nacon.numpy()[0], d.ncollision.numpy()[0]]))
nefc.append(np.max(d.nefc.numpy()))
if measure_solver_niter:
solver_niter.append(d.solver_niter.numpy())
@@ -150,7 +150,7 @@ def benchmark(
nsuccess = np.sum(~np.any(np.isnan(d.qpos.numpy()), axis=1))
run_duration = np.sum(time_vec)
return jit_duration, run_duration, trace, ncon, nefc, solver_niter, nsuccess
return jit_duration, run_duration, trace, nacon, nefc, solver_niter, nsuccess
class BenchmarkSuite:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+13 -13
View File
@@ -95,7 +95,7 @@ def ccd_kernel_builder(
opt_ccd_tolerance: wp.array(dtype=float),
geom_type: wp.array(dtype=int),
# Data in:
nconmax_in: int,
naconmax_in: int,
epa_vert_in: wp.array2d(dtype=wp.vec3),
epa_vert1_in: wp.array2d(dtype=wp.vec3),
epa_vert2_in: wp.array2d(dtype=wp.vec3),
@@ -135,7 +135,7 @@ def ccd_kernel_builder(
x2: wp.vec3,
count: int,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -222,7 +222,7 @@ def ccd_kernel_builder(
for i in range(ncontact):
write_contact(
nconmax_in,
naconmax_in,
dist,
points[i],
frame,
@@ -235,7 +235,7 @@ def ccd_kernel_builder(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -298,7 +298,7 @@ def ccd_kernel_builder(
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# Data in:
nconmax_in: int,
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_in: wp.array(dtype=wp.vec2i),
@@ -328,7 +328,7 @@ def ccd_kernel_builder(
multiccd_face1_in: wp.array2d(dtype=wp.vec3),
multiccd_face2_in: wp.array2d(dtype=wp.vec3),
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -503,7 +503,7 @@ def ccd_kernel_builder(
ncontact = eval_ccd_write_contact(
opt_ccd_tolerance,
geom_type,
nconmax_in,
naconmax_in,
epa_vert_in,
epa_vert1_in,
epa_vert2_in,
@@ -541,7 +541,7 @@ def ccd_kernel_builder(
x1,
geom2.pos,
count,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -561,7 +561,7 @@ def ccd_kernel_builder(
eval_ccd_write_contact(
opt_ccd_tolerance,
geom_type,
nconmax_in,
naconmax_in,
epa_vert_in,
epa_vert1_in,
epa_vert2_in,
@@ -599,7 +599,7 @@ def ccd_kernel_builder(
geom1.pos,
geom2.pos,
0,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -638,7 +638,7 @@ def convex_narrowphase(m: Model, d: Data):
if m.geom_pair_type_count[upper_trid_index(len(GeomType), g1, g2)]:
wp.launch(
ccd_kernel_builder(m.opt.legacy_gjk, g1, g2, m.opt.ccd_iterations, True, 1e9, g1 == GeomType.HFIELD),
dim=d.nconmax,
dim=d.naconmax,
inputs=[
m.opt.ccd_tolerance,
m.geom_type,
@@ -680,7 +680,7 @@ def convex_narrowphase(m: Model, d: Data):
m.pair_margin,
m.pair_gap,
m.pair_friction,
d.nconmax,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.collision_pair,
@@ -711,7 +711,7 @@ def convex_narrowphase(m: Model, d: Data):
d.multiccd_face2,
],
outputs=[
d.ncon,
d.nacon,
d.contact.dist,
d.contact.pos,
d.contact.frame,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+15 -15
View File
@@ -36,13 +36,13 @@ wp.set_module_options({"enable_backward": False})
@wp.kernel
def _zero_ncon_ncollision(
def _zero_nacon_ncollision(
# Data out:
ncon_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int),
):
ncollision_out[0] = 0
ncon_out[0] = 0
nacon_out[0] = 0
@wp.func
@@ -292,7 +292,7 @@ def _add_geom_pair(
geom_type: wp.array(dtype=int),
nxn_pairid: wp.array(dtype=int),
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
geom1: int,
geom2: int,
@@ -306,7 +306,7 @@ def _add_geom_pair(
):
pairid = wp.atomic_add(ncollision_out, 0, 1)
if pairid >= nconmax_in:
if pairid >= naconmax_in:
return
type1 = geom_type[geom1]
@@ -407,7 +407,7 @@ def _sap_broadphase(broadphase_filter):
nxn_pairid: wp.array(dtype=int),
# Data in:
nworld_in: int,
nconmax_in: int,
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore
@@ -455,7 +455,7 @@ def _sap_broadphase(broadphase_filter):
_add_geom_pair(
geom_type,
nxn_pairid,
nconmax_in,
naconmax_in,
geom1,
geom2,
worldid,
@@ -581,7 +581,7 @@ def sap_broadphase(m: Model, d: Data):
m.geom_margin,
m.nxn_pairid,
d.nworld,
d.nconmax,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.sap_sort_index.reshape((-1, m.ngeom)),
@@ -609,7 +609,7 @@ def _nxn_broadphase(broadphase_filter):
nxn_geom_pair: wp.array(dtype=wp.vec2i),
nxn_pairid: wp.array(dtype=int),
# Data in:
nconmax_in: int,
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
# Data out:
@@ -628,7 +628,7 @@ def _nxn_broadphase(broadphase_filter):
_add_geom_pair(
geom_type,
nxn_pairid,
nconmax_in,
naconmax_in,
geom1,
geom2,
worldid,
@@ -668,7 +668,7 @@ def nxn_broadphase(m: Model, d: Data):
m.geom_margin,
m.nxn_geom_pair_filtered,
m.nxn_pairid_filtered,
d.nconmax,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
],
@@ -702,18 +702,18 @@ def collision(m: Model, d: Data):
distance, position, and frame.
The results are used to populate the `d.contact` array, and the total number of contacts
is stored in `d.ncon`. If `d.ncon` is larger than `d.nconmax` then an overflow has
is stored in `d.nacon`. If `d.nacon` is larger than `d.naconmax` then an overflow has
occurred and the remaining contacts will be skipped. If this happens, raise the `nconmax`
parameter in `io.make_data` or `io.put_data`.
This function will do nothing except zero out arrays if collision detection is disabled
via `m.opt.disableflags` or if `d.nconmax` is 0.
via `m.opt.disableflags` or if `d.nacon` is 0.
"""
# zero contact and collision counters
wp.launch(_zero_ncon_ncollision, dim=1, outputs=[d.ncon, d.ncollision])
wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
if d.nconmax == 0 or m.opt.disableflags & (DisableBit.CONSTRAINT | DisableBit.CONTACT):
if d.naconmax == 0 or m.opt.disableflags & (DisableBit.CONSTRAINT | DisableBit.CONTACT):
return
if m.opt.broadphase == BroadphaseType.NXN:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
Vendored
+88 -84
View File
@@ -170,26 +170,22 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
# exhaustive search over all vertices
if convex.graphadr == -1 or convex.vertnum < 10:
# find support points
max_support = wp.float32(-_HUGE_VAL)
for i in range(convex.vertnum):
support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + i], n)
max_support = wp.max(support, max_support)
threshold = wp.max(0.0, max_support - 1e-3)
# find first support point (a)
a_dist = wp.float32(-_HUGE_VAL)
max_support = wp.float32(-_HUGE_VAL)
a = wp.vec3()
for i in range(convex.vertnum):
vert = convex.vert[convex.vertadr + i]
support = wp.dot(plane_pos_local - vert, n)
dist = wp.where(support > threshold, support, -_HUGE_VAL)
if dist > a_dist:
if support > max_support:
max_support = support
indices[0] = i
a_dist = dist
a = vert
if max_support < 0:
return contact_dist, contact_pos, plane_normal
threshold = max_support - 1e-3
# find point (b) furthest from a
b_dist = wp.float32(-_HUGE_VAL)
b = wp.vec3()
@@ -377,7 +373,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
@wp.func
def write_contact(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
dist_in: float,
pos_in: wp.vec3,
@@ -392,7 +388,7 @@ def write_contact(
geoms_in: wp.vec2i,
worldid_in: int,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -406,8 +402,8 @@ def write_contact(
contact_worldid_out: wp.array(dtype=int),
):
if dist_in - margin_in < 0.0:
cid = wp.atomic_add(ncon_out, 0, 1)
if cid < nconmax_in:
cid = wp.atomic_add(nacon_out, 0, 1)
if cid < naconmax_in:
contact_dist_out[cid] = dist_in
contact_pos_out[cid] = pos_in
contact_frame_out[cid] = frame_in
@@ -515,7 +511,7 @@ def contact_params(
@wp.func
def plane_sphere_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
plane: Geom,
sphere: Geom,
@@ -529,7 +525,7 @@ def plane_sphere_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -547,7 +543,7 @@ def plane_sphere_wrapper(
if dist - margin < 0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(plane.normal),
@@ -560,7 +556,7 @@ def plane_sphere_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -578,7 +574,7 @@ def plane_sphere_wrapper(
@wp.func
def sphere_sphere_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
sphere1: Geom,
sphere2: Geom,
@@ -592,7 +588,7 @@ def sphere_sphere_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -610,7 +606,7 @@ def sphere_sphere_wrapper(
if dist - margin < 0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(normal),
@@ -623,7 +619,7 @@ def sphere_sphere_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -641,7 +637,7 @@ def sphere_sphere_wrapper(
@wp.func
def sphere_capsule_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
sphere: Geom,
cap: Geom,
@@ -655,7 +651,7 @@ def sphere_capsule_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -676,7 +672,7 @@ def sphere_capsule_wrapper(
if dist - margin < 0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(normal),
@@ -689,7 +685,7 @@ def sphere_capsule_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -707,7 +703,7 @@ def sphere_capsule_wrapper(
@wp.func
def capsule_capsule_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
cap1: Geom,
cap2: Geom,
@@ -721,7 +717,7 @@ def capsule_capsule_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -752,7 +748,7 @@ def capsule_capsule_wrapper(
if dist - margin < 0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(normal),
@@ -765,7 +761,7 @@ def capsule_capsule_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -783,7 +779,7 @@ def capsule_capsule_wrapper(
@wp.func
def plane_capsule_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
plane: Geom,
cap: Geom,
@@ -797,7 +793,7 @@ def plane_capsule_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -827,7 +823,7 @@ def plane_capsule_wrapper(
disti = dist[i]
if disti - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
disti,
pos[i],
frame,
@@ -840,7 +836,7 @@ def plane_capsule_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -858,7 +854,7 @@ def plane_capsule_wrapper(
@wp.func
def plane_ellipsoid_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
plane: Geom,
ellipsoid: Geom,
@@ -872,7 +868,7 @@ def plane_ellipsoid_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -890,7 +886,7 @@ def plane_ellipsoid_wrapper(
if dist - margin < 0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(normal),
@@ -903,7 +899,7 @@ def plane_ellipsoid_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -921,7 +917,7 @@ def plane_ellipsoid_wrapper(
@wp.func
def plane_box_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
plane: Geom,
box: Geom,
@@ -935,7 +931,7 @@ def plane_box_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -956,7 +952,7 @@ def plane_box_wrapper(
disti = dist[i]
if disti - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
disti,
pos[i],
frame,
@@ -969,7 +965,7 @@ def plane_box_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -990,7 +986,7 @@ _HUGE_VAL = 1e6
@wp.func
def plane_convex_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
plane: Geom,
convex: Geom,
@@ -1004,7 +1000,7 @@ def plane_convex_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1025,7 +1021,7 @@ def plane_convex_wrapper(
disti = dist[i]
if disti - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
disti,
pos[i],
frame,
@@ -1038,7 +1034,7 @@ def plane_convex_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1056,7 +1052,7 @@ def plane_convex_wrapper(
@wp.func
def sphere_cylinder_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
sphere: Geom,
cylinder: Geom,
@@ -1070,7 +1066,7 @@ def sphere_cylinder_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1098,7 +1094,7 @@ def sphere_cylinder_wrapper(
if dist - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(normal),
@@ -1111,7 +1107,7 @@ def sphere_cylinder_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1129,7 +1125,7 @@ def sphere_cylinder_wrapper(
@wp.func
def plane_cylinder_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
plane: Geom,
cylinder: Geom,
@@ -1143,7 +1139,7 @@ def plane_cylinder_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1174,7 +1170,7 @@ def plane_cylinder_wrapper(
disti = dist[i]
if disti - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
disti,
pos[i],
frame,
@@ -1187,7 +1183,7 @@ def plane_cylinder_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1205,7 +1201,7 @@ def plane_cylinder_wrapper(
@wp.func
def sphere_box_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
sphere: Geom,
box: Geom,
@@ -1219,7 +1215,7 @@ def sphere_box_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1236,7 +1232,7 @@ def sphere_box_wrapper(
if dist - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(normal),
@@ -1249,7 +1245,7 @@ def sphere_box_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1267,7 +1263,7 @@ def sphere_box_wrapper(
@wp.func
def capsule_box_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
cap: Geom,
box: Geom,
@@ -1281,7 +1277,7 @@ def capsule_box_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1314,7 +1310,7 @@ def capsule_box_wrapper(
disti = dist[i]
if disti - margin < 0.0:
write_contact(
nconmax_in,
naconmax_in,
disti,
pos[i],
make_frame(normal[i]),
@@ -1327,7 +1323,7 @@ def capsule_box_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1345,7 +1341,7 @@ def capsule_box_wrapper(
@wp.func
def box_box_wrapper(
# Data in:
nconmax_in: int,
naconmax_in: int,
# In:
box1: Geom,
box2: Geom,
@@ -1359,7 +1355,7 @@ def box_box_wrapper(
solimp: vec5,
geoms: wp.vec2i,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1388,7 +1384,7 @@ def box_box_wrapper(
continue
write_contact(
nconmax_in,
naconmax_in,
dist[i],
pos[i],
make_frame(normal[i]),
@@ -1401,7 +1397,7 @@ def box_box_wrapper(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1446,12 +1442,17 @@ def _check_primitive_collisions():
assert _check_primitive_collisions(), "_PRIMITIVE_COLLISIONS is in invalid order"
_primitive_collisions_types = []
_primitive_collisions_func = []
@cache_kernel
def _create_narrowphase_kernel(primitive_collisions_types, primitive_collisions_func):
# AD: no unique here:
# * we expect this generator to be called only once per model, so no repeated compilation
# * module="unique" is generating problems because it uses the function name as the key
# that in turn will cause multiple kernels to be generated with the same name
# this is mostly problematic in cases like the UTs where we don't clear the kernel cache
# between different tests.
def _create_narrowphase_kernel():
@nested_kernel(module="unique", enable_backward=False)
@nested_kernel(enable_backward=False)
def _primitive_narrowphase(
# Model:
geom_type: wp.array(dtype=int),
@@ -1492,7 +1493,7 @@ def _create_narrowphase_kernel():
pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5),
# Data in:
nconmax_in: int,
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_in: wp.array(dtype=wp.vec2i),
@@ -1500,7 +1501,7 @@ def _create_narrowphase_kernel():
collision_worldid_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int),
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -1595,13 +1596,13 @@ def _create_narrowphase_kernel():
geom_xmat_in[worldid, g2],
)
for i in range(wp.static(len(_primitive_collisions_func))):
collision_type1 = wp.static(_primitive_collisions_types[i][0])
collision_type2 = wp.static(_primitive_collisions_types[i][1])
for i in range(wp.static(len(primitive_collisions_func))):
collision_type1 = wp.static(primitive_collisions_types[i][0])
collision_type2 = wp.static(primitive_collisions_types[i][1])
if collision_type1 == type1 and collision_type2 == type2:
wp.static(_primitive_collisions_func[i])(
nconmax_in,
wp.static(primitive_collisions_func[i])(
naconmax_in,
geom1,
geom2,
worldid,
@@ -1613,7 +1614,7 @@ def _create_narrowphase_kernel():
solreffriction,
solimp,
geoms,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -1631,13 +1632,16 @@ def _create_narrowphase_kernel():
def _primitive_narrowphase_builder(m: Model):
_primitive_collisions_types = []
_primitive_collisions_func = []
for types, func in _PRIMITIVE_COLLISIONS.items():
idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
if m.geom_pair_type_count[idx] and types not in _primitive_collisions_types:
_primitive_collisions_types.append(types)
_primitive_collisions_func.append(func)
return _create_narrowphase_kernel()
return _create_narrowphase_kernel(_primitive_collisions_types, _primitive_collisions_func)
@event_scope
@@ -1660,7 +1664,7 @@ def primitive_narrowphase(m: Model, d: Data):
# for pair types without collisions, as well as updating the launch dimensions.
wp.launch(
_primitive_narrowphase_builder(m),
dim=d.nconmax,
dim=d.naconmax,
inputs=[
m.geom_type,
m.geom_condim,
@@ -1699,7 +1703,7 @@ def primitive_narrowphase(m: Model, d: Data):
m.pair_margin,
m.pair_gap,
m.pair_friction,
d.nconmax,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.collision_pair,
@@ -1708,7 +1712,7 @@ def primitive_narrowphase(m: Model, d: Data):
d.ncollision,
],
outputs=[
d.ncon,
d.nacon,
d.contact.dist,
d.contact.pos,
d.contact.frame,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+7 -7
View File
@@ -668,7 +668,7 @@ def _sdf_narrowphase(
plugin_attr: wp.array(dtype=wp.vec3f),
geom_plugin_index: wp.array(dtype=int),
# Data in:
nconmax_in: int,
naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_in: wp.array(dtype=wp.vec2i),
@@ -679,7 +679,7 @@ def _sdf_narrowphase(
sdf_initpoints: int,
sdf_iterations: int,
# Data out:
ncon_out: wp.array(dtype=int),
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),
@@ -850,7 +850,7 @@ def _sdf_narrowphase(
mesh_data2,
)
write_contact(
nconmax_in,
naconmax_in,
dist,
pos,
make_frame(n),
@@ -863,7 +863,7 @@ def _sdf_narrowphase(
solimp,
geoms,
worldid,
ncon_out,
nacon_out,
contact_dist_out,
contact_pos_out,
contact_frame_out,
@@ -882,7 +882,7 @@ def _sdf_narrowphase(
def sdf_narrowphase(m: Model, d: Data):
wp.launch(
_sdf_narrowphase,
dim=(m.opt.sdf_initpoints, d.nconmax),
dim=(m.opt.sdf_initpoints, d.naconmax),
inputs=[
m.nmeshface,
m.geom_type,
@@ -931,7 +931,7 @@ def sdf_narrowphase(m: Model, d: Data):
m.plugin,
m.plugin_attr,
m.geom_plugin_index,
d.nconmax,
d.naconmax,
d.geom_xpos,
d.geom_xmat,
d.collision_pair,
@@ -942,7 +942,7 @@ def sdf_narrowphase(m: Model, d: Data):
m.opt.sdf_iterations,
],
outputs=[
d.ncon,
d.nacon,
d.contact.dist,
d.contact.pos,
d.contact.frame,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+8 -8
View File
@@ -1110,7 +1110,7 @@ def _efc_contact_pyramidal(
geom_bodyid: wp.array(dtype=int),
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
qvel_in: wp.array2d(dtype=float),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector),
@@ -1141,7 +1141,7 @@ def _efc_contact_pyramidal(
):
conid, dimid = wp.tid()
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
condim = condim_in[conid]
@@ -1275,7 +1275,7 @@ def _efc_contact_elliptic(
geom_bodyid: wp.array(dtype=int),
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
qvel_in: wp.array2d(dtype=float),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector),
@@ -1307,7 +1307,7 @@ def _efc_contact_elliptic(
):
conid, dimid = wp.tid()
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
condim = condim_in[conid]
@@ -1799,7 +1799,7 @@ def make_constraint(m: types.Model, d: types.Data):
if m.opt.cone == types.ConeType.PYRAMIDAL:
wp.launch(
_efc_contact_pyramidal,
dim=(d.nconmax, 2 * (m.condim_max - 1) if m.condim_max > 1 else 1),
dim=(d.naconmax, 2 * (m.condim_max - 1) if m.condim_max > 1 else 1),
inputs=[
m.nv,
m.opt.timestep,
@@ -1810,7 +1810,7 @@ def make_constraint(m: types.Model, d: types.Data):
m.dof_bodyid,
m.geom_bodyid,
d.njmax,
d.ncon,
d.nacon,
d.qvel,
d.subtree_com,
d.cdof,
@@ -1843,7 +1843,7 @@ def make_constraint(m: types.Model, d: types.Data):
elif m.opt.cone == types.ConeType.ELLIPTIC:
wp.launch(
_efc_contact_elliptic,
dim=(d.nconmax, m.condim_max),
dim=(d.naconmax, m.condim_max),
inputs=[
m.nv,
m.opt.timestep,
@@ -1854,7 +1854,7 @@ def make_constraint(m: types.Model, d: types.Data):
m.dof_bodyid,
m.geom_bodyid,
d.njmax,
d.ncon,
d.nacon,
d.qvel,
d.subtree_com,
d.cdof,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+12 -8
View File
@@ -206,9 +206,10 @@ def _next_time(
# Model:
opt_timestep: wp.array(dtype=float),
# Data in:
nconmax_in: int,
nworld_in: int,
naconmax_in: int,
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int),
time_in: wp.array(dtype=float),
ncollision_in: wp.array(dtype=int),
@@ -224,11 +225,13 @@ def _next_time(
if worldid == 0:
ncollision = ncollision_in[0]
if ncollision > nconmax_in:
wp.printf("ncollision overflow - please increase nconmax to %u\n", ncollision)
if ncollision > naconmax_in:
nconmax = int(wp.ceil(float(ncollision) / float(nworld_in)))
wp.printf("ncollision overflow - please increase nconmax to %u\n", nconmax)
if ncon_in[0] > nconmax_in:
wp.printf("ncon overflow - please increase nconmax to %u\n", ncon_in[0])
if nacon_in[0] > naconmax_in:
nconmax = int(wp.ceil(float(ncollision) / float(nworld_in)))
wp.printf("nacon overflow - please increase nconmax to %u\n", nconmax)
def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None):
@@ -299,9 +302,10 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
dim=(d.nworld,),
inputs=[
m.opt.timestep,
d.nconmax,
d.nworld,
d.naconmax,
d.njmax,
d.ncon,
d.nacon,
d.nefc,
d.time,
d.ncollision,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+464 -191
View File
@@ -109,9 +109,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any():
raise NotImplementedError("Flex collisions are not implemented.")
if mjm.geom_fluid.any():
raise NotImplementedError("Ellipsoid fluid model not implemented.")
# check options
for opt, opt_types, msg in (
(mjm.opt.integrator, types.IntegratorType, "Integrator"),
@@ -456,6 +453,17 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
"Collision sensors with height fields are not implemented.",
)
if mjm.geom_fluid.size:
geom_fluid_params = mjm.geom_fluid.reshape(mjm.ngeom, mujoco.mjNFLUID)
else:
geom_fluid_params = np.zeros((mjm.ngeom, mujoco.mjNFLUID))
body_fluid_ellipsoid = np.zeros(mjm.nbody, dtype=bool)
if mjm.ngeom:
active_geom = geom_fluid_params[:, 0] > 0
if np.any(active_geom):
body_fluid_ellipsoid[mjm.geom_bodyid[active_geom]] = True
m = types.Model(
nq=mjm.nq,
nv=mjm.nv,
@@ -566,6 +574,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
body_contype=wp.array(mjm.body_contype, dtype=int),
body_conaffinity=wp.array(mjm.body_conaffinity, dtype=int),
body_gravcomp=create_nmodel_batched_array(mjm.body_gravcomp, dtype=float),
body_fluid_ellipsoid=wp.array(body_fluid_ellipsoid, dtype=bool),
jnt_type=wp.array(mjm.jnt_type, dtype=int),
jnt_qposadr=wp.array(mjm.jnt_qposadr, dtype=int),
jnt_dofadr=wp.array(mjm.jnt_dofadr, dtype=int),
@@ -617,6 +626,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
geom_solref=create_nmodel_batched_array(mjm.geom_solref, dtype=wp.vec2),
geom_solimp=create_nmodel_batched_array(mjm.geom_solimp, dtype=types.vec5),
geom_size=create_nmodel_batched_array(mjm.geom_size, dtype=wp.vec3),
geom_fluid=wp.array(geom_fluid_params, dtype=float),
geom_aabb=wp.array2d(mjm.geom_aabb, dtype=wp.vec3),
geom_rbound=create_nmodel_batched_array(mjm.geom_rbound, dtype=float),
geom_pos=create_nmodel_batched_array(mjm.geom_pos, dtype=wp.vec3),
@@ -902,43 +912,48 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
return m
def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: int = -1) -> types.Data:
def make_data(
mjm: mujoco.MjModel,
nworld: int = 1,
nconmax: Optional[int] = None,
njmax: Optional[int] = None,
naconmax: Optional[int] = None,
) -> types.Data:
"""
Creates a data object on device.
Args:
mjm (mujoco.MjModel): The model containing kinematic and dynamic information (host).
nworld (int, optional): Number of worlds. Defaults to 1.
nconmax (int, optional): Maximum number of contacts for all worlds. Defaults to -1.
njmax (int, optional): Maximum number of constraints per world. Defaults to -1.
nworld (int, optional): The number of worlds. Defaults to 1.
nconmax (int, optional): Number of contacts to allocate per world. Contacts exist in large
heterogenous arrays: one world may have more than nconmax contacts.
njmax (int, optional): Number of constraints to allocate per world. Constraint arrays are
batched by world: no world may have more than njmax constraints.
naconmax (int, optional): Number of contacts to allocate for all worlds. Overrides nconmax.
Returns:
Data: The data object containing the current state and output arrays (device).
"""
# TODO(team): move to Model?
if nconmax == -1:
# TODO(team): heuristic for nconmax
nconmax = nworld * 20
if njmax == -1:
# TODO(team): heuristic for njmax
njmax = 20 * 6
# TODO(team): move nconmax, njmax to Model?
# TODO(team): improve heuristic for nconmax and njmax
nconmax = nconmax or 20
njmax = njmax or nconmax * 6
if nworld < 1 or nworld > MAX_WORLDS:
raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}")
if nconmax < 0:
raise ValueError("nconmax must be >= 0")
if naconmax is None:
if nconmax < 0:
raise ValueError("nconmax must be >= 0")
naconmax = max(512, nworld * nconmax)
elif naconmax < 0:
raise ValueError("naconmax must be >= 0")
if njmax < 0:
raise ValueError("njmax must be >= 0")
condim = np.concatenate((mjm.geom_condim, mjm.pair_dim))
condim_max = np.max(condim) if len(condim) > 0 else 0
max_npolygon = _max_npolygon(mjm)
max_meshdegree = _max_meshdegree(mjm)
if mujoco.mj_isSparse(mjm):
qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
@@ -950,15 +965,19 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
qM_integration = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
qLD_integration = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
condim = np.concatenate((mjm.geom_condim, mjm.pair_dim))
condim_max = np.max(condim) if len(condim) > 0 else 0
max_npolygon = _max_npolygon(mjm)
max_meshdegree = _max_meshdegree(mjm)
nsensorcontact = np.sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT)
nrangefinder = sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER)
return types.Data(
nworld=nworld,
nconmax=nconmax,
naconmax=naconmax,
njmax=njmax,
solver_niter=wp.zeros(nworld, dtype=int),
ncon=wp.zeros(1, dtype=int),
nacon=wp.zeros(1, dtype=int),
ne=wp.zeros(nworld, dtype=int),
ne_connect=wp.zeros(nworld, dtype=int), # warp only
ne_weld=wp.zeros(nworld, dtype=int), # warp only
@@ -1032,21 +1051,18 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
qfrc_constraint=wp.zeros((nworld, mjm.nv), dtype=float),
qfrc_inverse=wp.zeros((nworld, mjm.nv), dtype=float),
contact=types.Contact(
dist=wp.zeros((nconmax,), dtype=float),
pos=wp.zeros((nconmax,), dtype=wp.vec3f),
frame=wp.zeros((nconmax,), dtype=wp.mat33f),
includemargin=wp.zeros((nconmax,), dtype=float),
friction=wp.zeros((nconmax,), dtype=types.vec5),
solref=wp.zeros((nconmax,), dtype=wp.vec2f),
solreffriction=wp.zeros((nconmax,), dtype=wp.vec2f),
solimp=wp.zeros((nconmax,), dtype=types.vec5),
dim=wp.zeros((nconmax,), dtype=int),
geom=wp.zeros((nconmax,), dtype=wp.vec2i),
efc_address=wp.zeros(
(nconmax, np.maximum(1, 2 * (condim_max - 1))),
dtype=int,
),
worldid=wp.zeros((nconmax,), dtype=int),
dist=wp.zeros((naconmax,), dtype=float),
pos=wp.zeros((naconmax,), dtype=wp.vec3f),
frame=wp.zeros((naconmax,), dtype=wp.mat33f),
includemargin=wp.zeros((naconmax,), dtype=float),
friction=wp.zeros((naconmax,), dtype=types.vec5),
solref=wp.zeros((naconmax,), dtype=wp.vec2f),
solreffriction=wp.zeros((naconmax,), dtype=wp.vec2f),
solimp=wp.zeros((naconmax,), dtype=types.vec5),
dim=wp.zeros((naconmax,), dtype=int),
geom=wp.zeros((naconmax,), dtype=wp.vec2i),
efc_address=wp.zeros((naconmax, np.maximum(1, 2 * (condim_max - 1))), dtype=int),
worldid=wp.zeros((naconmax,), dtype=int),
),
efc=types.Constraint(
type=wp.zeros((nworld, njmax), dtype=int),
@@ -1109,33 +1125,33 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
np.array([i * mjm.ngeom if i < nworld + 1 else 0 for i in range(2 * nworld)]).reshape((nworld, 2)), dtype=int
),
# collision driver
collision_pair=wp.zeros((nconmax,), dtype=wp.vec2i),
collision_pairid=wp.zeros((nconmax,), dtype=int),
collision_worldid=wp.zeros((nconmax,), dtype=int),
collision_pair=wp.zeros((naconmax,), dtype=wp.vec2i),
collision_pairid=wp.zeros((naconmax,), dtype=int),
collision_worldid=wp.zeros((naconmax,), dtype=int),
ncollision=wp.zeros((1,), dtype=int),
# narrowphase (EPA polytope)
epa_vert=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert1=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert2=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert_index1=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_vert_index2=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_face=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3i),
epa_pr=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_norm2=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=float),
epa_index=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_map=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_horizon=wp.zeros(shape=(nconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int),
multiccd_polygon=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_clipped=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_pnormal=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3),
multiccd_pdist=wp.zeros(shape=(nconmax, max_npolygon), dtype=float),
multiccd_idx1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int),
multiccd_idx2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int),
multiccd_n1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3),
multiccd_n2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3),
multiccd_endvert=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3),
multiccd_face1=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3),
multiccd_face2=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3),
epa_vert=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert1=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert2=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert_index1=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_vert_index2=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_face=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3i),
epa_pr=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_norm2=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=float),
epa_index=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_map=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_horizon=wp.zeros(shape=(naconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int),
multiccd_polygon=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_clipped=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_pnormal=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_pdist=wp.zeros(shape=(naconmax, max_npolygon), dtype=float),
multiccd_idx1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_idx2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_n1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_n2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_endvert=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_face1=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_face2=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
# rne_postconstraint
cacc=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector),
cfrc_int=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector),
@@ -1176,9 +1192,10 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
def put_data(
mjm: mujoco.MjModel,
mjd: mujoco.MjData,
nworld: Optional[int] = None,
nworld: int = 1,
nconmax: Optional[int] = None,
njmax: Optional[int] = None,
naconmax: Optional[int] = None,
) -> types.Data:
"""
Moves data from host to a device.
@@ -1187,8 +1204,11 @@ def put_data(
mjm (mujoco.MjModel): The model containing kinematic and dynamic information (host).
mjd (mujoco.MjData): The data object containing current state and output arrays (host).
nworld (int, optional): The number of worlds. Defaults to 1.
nconmax (int, optional): The maximum number of contacts for all worlds. Defaults to -1.
njmax (int, optional): The maximum number of constraints per world. Defaults to -1.
nconmax (int, optional): Number of contacts to allocate per world. Contacts exist in large
heterogenous arrays: one world may have more than nconmax contacts.
njmax (int, optional): Number of constraints to allocate per world. Constraint arrays are
batched by world: no world may have more than njmax constraints.
naconmax (int, optional): Number of contacts to allocate for all worlds. Overrides nconmax.
Returns:
Data: The data object containing the current state and output arrays (device).
@@ -1197,24 +1217,27 @@ def put_data(
# TODO(team): decide what to do about uninitialized warp-only fields created by put_data
# we need to ensure these are only workspace fields and don't carry state
nworld = nworld or 1
# TODO(team): better heuristic for nconmax
nconmax = nconmax or max(512, 4 * mjd.ncon * nworld)
# TODO(team): better heuristic for njmax
# TODO(team): better heuristic for nconmax and njmax
nconmax = nconmax or max(5, 4 * mjd.ncon)
njmax = njmax or max(5, 4 * mjd.nefc)
if nworld < 1 or nworld > MAX_WORLDS:
raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}")
if nconmax < 0:
raise ValueError("nconmax must be >= 0")
if naconmax is None:
if nconmax < 0:
raise ValueError("nconmax must be >= 0")
if mjd.ncon > nconmax:
raise ValueError(f"nconmax overflow (nconmax must be >= {mjd.ncon})")
naconmax = max(512, nworld * nconmax)
elif naconmax < mjd.ncon * nworld:
raise ValueError(f"naconmax overflow (naconmax must be >= {mjd.ncon * nworld})")
if njmax < 0:
raise ValueError("njmax must be >= 0")
if nworld * mjd.ncon > nconmax:
raise ValueError(f"nconmax overflow (nconmax must be >= {nworld * mjd.ncon})")
if mjd.nefc > njmax:
raise ValueError(f"njmax overflow (njmax must be >= {mjd.nefc})")
@@ -1261,7 +1284,7 @@ def put_data(
condim = np.concatenate((mjm.geom_condim, mjm.pair_dim))
condim_max = np.max(condim) if len(condim) > 0 else 0
contact_efc_address = np.zeros((nconmax, np.maximum(1, 2 * (condim_max - 1))), dtype=int)
contact_efc_address = np.zeros((naconmax, np.maximum(1, 2 * (condim_max - 1))), dtype=int)
for i in range(nworld):
for j in range(mjd.ncon):
condim = mjd.contact.dim[j]
@@ -1275,7 +1298,7 @@ def put_data(
for k in range(nconvar):
contact_efc_address[i * mjd.ncon + j, k] = mjd.nefc * i + efc_address + k
contact_worldid = np.pad(np.repeat(np.arange(nworld), mjd.ncon), (0, nconmax - nworld * mjd.ncon))
contact_worldid = np.pad(np.repeat(np.arange(nworld), mjd.ncon), (0, naconmax - nworld * mjd.ncon))
ne_connect = int(3 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_CONNECT) & mjd.eq_active))
ne_weld = int(6 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_WELD) & mjd.eq_active))
@@ -1335,10 +1358,10 @@ def put_data(
return types.Data(
nworld=nworld,
nconmax=nconmax,
naconmax=naconmax,
njmax=njmax,
solver_niter=tile(mjd.solver_niter[0]),
ncon=arr([mjd.ncon * nworld]),
nacon=arr([mjd.ncon * nworld]),
ne=wp.full(shape=(nworld), value=mjd.ne),
ne_connect=wp.full(shape=(nworld), value=ne_connect),
ne_weld=wp.full(shape=(nworld), value=ne_weld),
@@ -1412,16 +1435,16 @@ def put_data(
qfrc_constraint=tile(mjd.qfrc_constraint),
qfrc_inverse=tile(mjd.qfrc_inverse),
contact=types.Contact(
dist=padtile(mjd.contact.dist, nconmax),
pos=padtile(mjd.contact.pos, nconmax, dtype=wp.vec3),
frame=padtile(mjd.contact.frame, nconmax, dtype=wp.mat33),
includemargin=padtile(mjd.contact.includemargin, nconmax),
friction=padtile(mjd.contact.friction, nconmax, dtype=types.vec5),
solref=padtile(mjd.contact.solref, nconmax, dtype=wp.vec2f),
solreffriction=padtile(mjd.contact.solreffriction, nconmax, dtype=wp.vec2f),
solimp=padtile(mjd.contact.solimp, nconmax, dtype=types.vec5),
dim=padtile(mjd.contact.dim, nconmax),
geom=padtile(mjd.contact.geom, nconmax, dtype=wp.vec2i),
dist=padtile(mjd.contact.dist, naconmax),
pos=padtile(mjd.contact.pos, naconmax, dtype=wp.vec3),
frame=padtile(mjd.contact.frame, naconmax, dtype=wp.mat33),
includemargin=padtile(mjd.contact.includemargin, naconmax),
friction=padtile(mjd.contact.friction, naconmax, dtype=types.vec5),
solref=padtile(mjd.contact.solref, naconmax, dtype=wp.vec2f),
solreffriction=padtile(mjd.contact.solreffriction, naconmax, dtype=wp.vec2f),
solimp=padtile(mjd.contact.solimp, naconmax, dtype=types.vec5),
dim=padtile(mjd.contact.dim, naconmax),
geom=padtile(mjd.contact.geom, naconmax, dtype=wp.vec2i),
efc_address=arr(contact_efc_address),
worldid=arr(contact_worldid),
),
@@ -1483,33 +1506,33 @@ def put_data(
sap_cumulative_sum=wp.zeros((nworld, mjm.ngeom), dtype=int),
sap_segment_index=arr(np.array([i * mjm.ngeom if i < nworld + 1 else 0 for i in range(2 * nworld)]).reshape((nworld, 2))),
# collision driver
collision_pair=wp.empty(nconmax, dtype=wp.vec2i),
collision_pairid=wp.empty(nconmax, dtype=int),
collision_worldid=wp.empty(nconmax, dtype=int),
collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
collision_pairid=wp.empty(naconmax, dtype=int),
collision_worldid=wp.empty(naconmax, dtype=int),
ncollision=wp.zeros(1, dtype=int),
# narrowphase (EPA polytope)
epa_vert=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert1=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert2=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert_index1=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_vert_index2=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_face=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3i),
epa_pr=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_norm2=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=float),
epa_index=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_map=wp.zeros(shape=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_horizon=wp.zeros(shape=(nconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int),
multiccd_polygon=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_clipped=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_pnormal=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3),
multiccd_pdist=wp.zeros(shape=(nconmax, max_npolygon), dtype=float),
multiccd_idx1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int),
multiccd_idx2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int),
multiccd_n1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3),
multiccd_n2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3),
multiccd_endvert=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3),
multiccd_face1=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3),
multiccd_face2=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3),
epa_vert=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert1=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert2=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert_index1=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_vert_index2=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_face=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3i),
epa_pr=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_norm2=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=float),
epa_index=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_map=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_horizon=wp.zeros(shape=(naconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int),
multiccd_polygon=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_clipped=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_pnormal=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_pdist=wp.zeros(shape=(naconmax, max_npolygon), dtype=float),
multiccd_idx1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_idx2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_n1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_n2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_endvert=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_face1=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_face2=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
# rne_postconstraint but also smooth
cacc=tile(mjd.cacc, dtype=wp.spatial_vector),
cfrc_int=tile(mjd.cfrc_int, dtype=wp.spatial_vector),
@@ -1565,11 +1588,11 @@ def get_data_into(
result.solver_niter[0] = d.solver_niter.numpy()[0]
ncon = d.ncon.numpy()[0]
nacon = d.nacon.numpy()[0]
nefc = d.nefc.numpy()[0]
if ncon != result.ncon or nefc != result.nefc:
mujoco._functions._realloc_con_efc(result, ncon=ncon, nefc=nefc)
if nacon != result.ncon or nefc != result.nefc:
mujoco._functions._realloc_con_efc(result, ncon=nacon, nefc=nefc)
result.time = d.time.numpy()[0]
result.energy = d.energy.numpy()[0]
@@ -1635,16 +1658,16 @@ def get_data_into(
result.act = d.act.numpy()[0]
result.act_dot = d.act_dot.numpy()[0]
result.contact.dist[:] = d.contact.dist.numpy()[:ncon]
result.contact.pos[:] = d.contact.pos.numpy()[:ncon]
result.contact.frame[:] = d.contact.frame.numpy()[:ncon].reshape((-1, 9))
result.contact.includemargin[:] = d.contact.includemargin.numpy()[:ncon]
result.contact.friction[:] = d.contact.friction.numpy()[:ncon]
result.contact.solref[:] = d.contact.solref.numpy()[:ncon]
result.contact.solreffriction[:] = d.contact.solreffriction.numpy()[:ncon]
result.contact.solimp[:] = d.contact.solimp.numpy()[:ncon]
result.contact.dim[:] = d.contact.dim.numpy()[:ncon]
result.contact.efc_address[:] = d.contact.efc_address.numpy()[:ncon, 0]
result.contact.dist[:] = d.contact.dist.numpy()[:nacon]
result.contact.pos[:] = d.contact.pos.numpy()[:nacon]
result.contact.frame[:] = d.contact.frame.numpy()[:nacon].reshape((-1, 9))
result.contact.includemargin[:] = d.contact.includemargin.numpy()[:nacon]
result.contact.friction[:] = d.contact.friction.numpy()[:nacon]
result.contact.solref[:] = d.contact.solref.numpy()[:nacon]
result.contact.solreffriction[:] = d.contact.solreffriction.numpy()[:nacon]
result.contact.solimp[:] = d.contact.solimp.numpy()[:nacon]
result.contact.dim[:] = d.contact.dim.numpy()[:nacon]
result.contact.efc_address[:] = d.contact.efc_address.numpy()[:nacon, 0]
if mujoco.mj_isSparse(mjm):
result.qM[:] = d.qM.numpy()[0, 0]
@@ -1704,8 +1727,43 @@ def get_data_into(
result.sensordata[:] = d.sensordata.numpy()
# TODO(thowell): shared @wp.func for _reset kernel?
@wp.kernel
def _reset_nworld(
def _reset_xfrc_applied_all(xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector)):
worldid, bodyid, elemid = wp.tid()
xfrc_applied_out[worldid, bodyid][elemid] = 0.0
@wp.kernel
def _reset_xfrc_applied(reset_in: wp.array(dtype=bool), xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector)):
worldid, bodyid, elemid = wp.tid()
if not reset_in[worldid]:
return
xfrc_applied_out[worldid, bodyid][elemid] = 0.0
@wp.kernel
def _reset_qM_all(qM_out: wp.array3d(dtype=float)):
worldid, elemid1, elemid2 = wp.tid()
qM_out[worldid, elemid1, elemid2] = 0.0
@wp.kernel
def _reset_qM(reset_in: wp.array(dtype=bool), qM_out: wp.array3d(dtype=float)):
worldid, elemid1, elemid2 = wp.tid()
if not reset_in[worldid]:
return
qM_out[worldid, elemid1, elemid2] = 0.0
@wp.kernel
def _reset_nworld_all(
# Model:
nq: int,
nv: int,
@@ -1719,7 +1777,7 @@ def _reset_nworld(
nworld_in: int,
# Data out:
solver_niter_out: wp.array(dtype=int),
ncon_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
ne_out: wp.array(dtype=int),
ne_connect_out: wp.array(dtype=int),
ne_weld_out: wp.array(dtype=int),
@@ -1746,7 +1804,7 @@ def _reset_nworld(
solver_niter_out[worldid] = 0
if worldid == 0:
ncon_out[0] = 0
nacon_out[0] = 0
ne_out[worldid] = 0
ne_connect_out[worldid] = 0
ne_weld_out[worldid] = 0
@@ -1778,7 +1836,85 @@ def _reset_nworld(
@wp.kernel
def _reset_mocap(
def _reset_nworld(
# Model:
nq: int,
nv: int,
nu: int,
na: int,
neq: int,
nsensordata: int,
qpos0: wp.array2d(dtype=float),
eq_active0: wp.array(dtype=bool),
# Data in:
nworld_in: int,
# In:
reset_in: wp.array(dtype=bool),
# Data out:
solver_niter_out: wp.array(dtype=int),
nacon_out: wp.array(dtype=int),
ne_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),
nf_out: wp.array(dtype=int),
nl_out: wp.array(dtype=int),
nefc_out: wp.array(dtype=int),
nsolving_out: wp.array(dtype=int),
time_out: wp.array(dtype=float),
energy_out: wp.array(dtype=wp.vec2),
qpos_out: wp.array2d(dtype=float),
qvel_out: wp.array2d(dtype=float),
act_out: wp.array2d(dtype=float),
qacc_warmstart_out: wp.array2d(dtype=float),
ctrl_out: wp.array2d(dtype=float),
qfrc_applied_out: wp.array2d(dtype=float),
eq_active_out: wp.array2d(dtype=bool),
qacc_out: wp.array2d(dtype=float),
act_dot_out: wp.array2d(dtype=float),
sensordata_out: wp.array2d(dtype=float),
):
worldid = wp.tid()
if not reset_in[worldid]:
return
solver_niter_out[worldid] = 0
if worldid == 0:
nacon_out[0] = 0
ne_out[worldid] = 0
ne_connect_out[worldid] = 0
ne_weld_out[worldid] = 0
ne_jnt_out[worldid] = 0
ne_ten_out[worldid] = 0
nf_out[worldid] = 0
nl_out[worldid] = 0
nefc_out[worldid] = 0
if worldid == 0:
nsolving_out[0] = nworld_in
time_out[worldid] = 0.0
energy_out[worldid] = wp.vec2(0.0, 0.0)
for i in range(nq):
qpos_out[worldid, i] = qpos0[worldid, i]
if i < nv:
qvel_out[worldid, i] = 0.0
qacc_warmstart_out[worldid, i] = 0.0
qfrc_applied_out[worldid, i] = 0.0
qacc_out[worldid, i] = 0.0
for i in range(nu):
ctrl_out[worldid, i] = 0.0
if i < na:
act_out[worldid, i] = 0.0
act_dot_out[worldid, i] = 0.0
for i in range(neq):
eq_active_out[worldid, i] = eq_active0[i]
for i in range(nsensordata):
sensordata_out[worldid, i] = 0.0
@wp.kernel
def _reset_mocap_all(
# Model:
body_mocapid: wp.array(dtype=int),
body_pos: wp.array2d(dtype=wp.vec3),
@@ -1797,9 +1933,33 @@ def _reset_mocap(
@wp.kernel
def _reset_contact(
def _reset_mocap(
# Model:
body_mocapid: wp.array(dtype=int),
body_pos: wp.array2d(dtype=wp.vec3),
body_quat: wp.array2d(dtype=wp.quat),
# In:
reset_in: wp.array(dtype=bool),
# Data out:
mocap_pos_out: wp.array2d(dtype=wp.vec3),
mocap_quat_out: wp.array2d(dtype=wp.quat),
):
worldid, bodyid = wp.tid()
if not reset_in[worldid]:
return
mocapid = body_mocapid[bodyid]
if mocapid >= 0:
mocap_pos_out[worldid, mocapid] = body_pos[worldid, bodyid]
mocap_quat_out[worldid, mocapid] = body_quat[worldid, bodyid]
@wp.kernel
def _reset_contact_all(
# Data in:
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
# In:
nefcaddress: int,
# Data out:
@@ -1818,9 +1978,55 @@ def _reset_contact(
):
conid = wp.tid()
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
contact_dist_out[conid] = 0.0
contact_pos_out[conid] = wp.vec3(0.0, 0.0, 0.0)
contact_frame_out[conid] = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
contact_includemargin_out[conid] = 0.0
contact_friction_out[conid] = types.vec5(0.0, 0.0, 0.0, 0.0, 0.0)
contact_solref_out[conid] = wp.vec2(0.0, 0.0)
contact_solreffriction_out[conid] = wp.vec2(0.0, 0.0)
contact_solimp_out[conid] = types.vec5(0.0, 0.0, 0.0, 0.0, 0.0)
contact_dim_out[conid] = 0
contact_geom_out[conid] = wp.vec2i(0, 0)
for i in range(nefcaddress):
contact_efc_address_out[conid, i] = 0
contact_worldid_out[conid] = 0
@wp.kernel
def _reset_contact(
# Data in:
nacon_in: wp.array(dtype=int),
# In:
reset_in: wp.array(dtype=bool),
nefcaddress: int,
# Data out:
contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33),
contact_includemargin_out: wp.array(dtype=float),
contact_friction_out: wp.array(dtype=types.vec5),
contact_solref_out: wp.array(dtype=wp.vec2),
contact_solreffriction_out: wp.array(dtype=wp.vec2),
contact_solimp_out: wp.array(dtype=types.vec5),
contact_dim_out: wp.array(dtype=int),
contact_geom_out: wp.array(dtype=wp.vec2i),
contact_efc_address_out: wp.array2d(dtype=int),
contact_worldid_out: wp.array(dtype=int),
):
conid = wp.tid()
if conid >= nacon_in[0]:
return
worldid = contact_worldid_out[conid]
if worldid >= 0:
if not reset_in[worldid]:
return
contact_dist_out[conid] = 0.0
contact_pos_out[conid] = wp.vec3(0.0)
contact_frame_out[conid] = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
@@ -1836,67 +2042,134 @@ def _reset_contact(
contact_worldid_out[conid] = 0
def reset_data(m: types.Model, d: types.Data):
def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
"""Clear data, set defaults."""
d.xfrc_applied.zero_()
d.qM.zero_()
if m.opt.is_sparse:
qM_dim = (1, m.nM)
else:
qM_dim = (m.nv, m.nv)
# set mocap_pos/quat = body_pos/quat for mocap bodies
wp.launch(
_reset_mocap, dim=(d.nworld, m.nbody), inputs=[m.body_mocapid, m.body_pos, m.body_quat], outputs=[d.mocap_pos, d.mocap_quat]
)
if reset is not None:
wp.launch(_reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset], outputs=[d.xfrc_applied])
wp.launch(_reset_qM, dim=(d.nworld, qM_dim[0], qM_dim[1]), inputs=[reset], outputs=[d.qM])
# clear contacts
wp.launch(
_reset_contact,
dim=d.nconmax,
inputs=[d.ncon, d.contact.efc_address.shape[1]],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
],
)
# set mocap_pos/quat = body_pos/quat for mocap bodies
wp.launch(
_reset_mocap,
dim=(d.nworld, m.nbody),
inputs=[m.body_mocapid, m.body_pos, m.body_quat, reset],
outputs=[d.mocap_pos, d.mocap_quat],
)
wp.launch(
_reset_nworld,
dim=d.nworld,
inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld],
outputs=[
d.solver_niter,
d.ncon,
d.ne,
d.ne_connect,
d.ne_weld,
d.ne_jnt,
d.ne_ten,
d.nf,
d.nl,
d.nefc,
d.nsolving,
d.time,
d.energy,
d.qpos,
d.qvel,
d.act,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
d.eq_active,
d.qacc,
d.act_dot,
d.sensordata,
],
)
# clear contacts
wp.launch(
_reset_contact,
dim=d.naconmax,
inputs=[d.nacon, reset, d.contact.efc_address.shape[1]],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
],
)
wp.launch(
_reset_nworld,
dim=d.nworld,
inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld, reset],
outputs=[
d.solver_niter,
d.nacon,
d.ne,
d.ne_connect,
d.ne_weld,
d.ne_jnt,
d.ne_ten,
d.nf,
d.nl,
d.nefc,
d.nsolving,
d.time,
d.energy,
d.qpos,
d.qvel,
d.act,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
d.eq_active,
d.qacc,
d.act_dot,
d.sensordata,
],
)
else:
wp.launch(_reset_xfrc_applied_all, dim=(d.nworld, m.nbody, 6), outputs=[d.xfrc_applied])
wp.launch(_reset_qM_all, dim=(d.nworld, qM_dim[0], qM_dim[1]), outputs=[d.qM])
wp.launch(
_reset_mocap_all,
dim=(d.nworld, m.nbody),
inputs=[m.body_mocapid, m.body_pos, m.body_quat],
outputs=[d.mocap_pos, d.mocap_quat],
)
wp.launch(
_reset_contact_all,
dim=d.naconmax,
inputs=[d.nacon, d.contact.efc_address.shape[1]],
outputs=[
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.includemargin,
d.contact.friction,
d.contact.solref,
d.contact.solreffriction,
d.contact.solimp,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
],
)
wp.launch(
_reset_nworld_all,
dim=d.nworld,
inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld],
outputs=[
d.solver_niter,
d.nacon,
d.ne,
d.ne_connect,
d.ne_weld,
d.ne_jnt,
d.ne_ten,
d.nf,
d.nl,
d.nefc,
d.nsolving,
d.time,
d.energy,
d.qpos,
d.qvel,
d.act,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
d.eq_active,
d.qacc,
d.act_dot,
d.sensordata,
],
)
def override_model(model: Union[types.Model, mujoco.MjModel], overrides: Union[dict[str, Any], Sequence[str]]):
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+212 -39
View File
@@ -20,6 +20,7 @@ from mujoco.mjx.third_party.mujoco_warp._src import support
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 DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
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.warp_util import event_scope
@@ -27,6 +28,45 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
@wp.func
def _pow2(val: float) -> float:
return val * val
@wp.func
def _pow4(val: float) -> float:
sq = val * val
return sq * sq
@wp.func
def _geom_semiaxes(size: wp.vec3, geom_type: int) -> wp.vec3: # kernel_analyzer: ignore
if geom_type == GeomType.SPHERE:
r = size[0]
return wp.vec3(r, r, r)
if geom_type == GeomType.CAPSULE:
radius = size[0]
half_length = size[1]
return wp.vec3(radius, radius, half_length + radius)
if geom_type == GeomType.CYLINDER:
radius = size[0]
half_length = size[1]
return wp.vec3(radius, radius, half_length)
# ellipsoid, box, mesh, sdf -> use size directly
return size
@wp.func
def _ellipsoid_max_moment(size: wp.vec3, dir: int) -> float:
d0 = size[dir]
d1 = size[(dir + 1) % 3]
d2 = size[(dir + 2) % 3]
return wp.static(8.0 / 15.0 * wp.pi) * d0 * _pow4(wp.max(d1, d2))
@wp.kernel
def _spring_damper_dof_passive(
# Model:
@@ -224,50 +264,185 @@ def _gravity_force(
@wp.kernel
def _box_fluid(
def _fluid_force(
# Model:
opt_wind: wp.array(dtype=wp.vec3),
opt_density: wp.array(dtype=float),
opt_viscosity: wp.array(dtype=float),
body_rootid: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int),
body_geomadr: wp.array(dtype=int),
body_mass: wp.array2d(dtype=float),
body_inertia: wp.array2d(dtype=wp.vec3),
body_fluid_ellipsoid: wp.array(dtype=bool),
geom_type: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_fluid: wp.array2d(dtype=float),
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# Data out:
fluid_applied_out: wp.array2d(dtype=wp.spatial_vector),
):
"""Fluid forces based on inertia-box approximation."""
"""Computes body-space fluid forces for both inertia-box and ellipsoid models."""
worldid, bodyid = wp.tid()
zero_force = wp.spatial_vector(wp.vec3(0.0), wp.vec3(0.0))
if bodyid == 0:
fluid_applied_out[worldid, bodyid] = zero_force
return
wind = opt_wind[worldid]
density = opt_density[worldid]
viscosity = opt_viscosity[worldid]
# map from CoM-centered to local body-centered 6D velocity
# body-inertial
pos = xipos_in[worldid, bodyid]
# Body kinematics
xipos = xipos_in[worldid, bodyid]
rot = ximat_in[worldid, bodyid]
rotT = wp.transpose(rot)
# transform velocity
cvel = cvel_in[worldid, bodyid]
torque = wp.spatial_top(cvel)
force = wp.spatial_bottom(cvel)
subtree_com = subtree_com_in[worldid, body_rootid[bodyid]]
dif = pos - subtree_com
force -= wp.cross(dif, torque)
ang_global = wp.spatial_top(cvel)
lin_global = wp.spatial_bottom(cvel)
subtree_root = subtree_com_in[worldid, body_rootid[bodyid]]
lin_com = lin_global - wp.cross(xipos - subtree_root, ang_global)
lvel_torque = rotT @ torque
lvel_force = rotT @ force
if body_fluid_ellipsoid[bodyid]:
force_global = wp.vec3(0.0)
torque_global = wp.vec3(0.0)
start = body_geomadr[bodyid]
count = body_geomnum[bodyid]
for i in range(count):
geomid = start + i
coef = geom_fluid[geomid, 0]
if coef <= 0.0:
continue
size = geom_size[worldid, geomid]
semiaxes = _geom_semiaxes(size, geom_type[geomid])
geom_rot = geom_xmat_in[worldid, geomid]
geom_rotT = wp.transpose(geom_rot)
geom_pos = geom_xpos_in[worldid, geomid]
lin_point = lin_com + wp.cross(ang_global, geom_pos - xipos)
l_ang = geom_rotT @ ang_global
l_lin = geom_rotT @ lin_point
if wind[0] or wind[1] or wind[2]:
l_lin -= geom_rotT @ wind
lfrc_torque = wp.vec3(0.0)
lfrc_force = wp.vec3(0.0)
if density > 0.0:
# added-mass forces and torques
virtual_mass = wp.vec3(geom_fluid[geomid, 6], geom_fluid[geomid, 7], geom_fluid[geomid, 8])
virtual_inertia = wp.vec3(geom_fluid[geomid, 9], geom_fluid[geomid, 10], geom_fluid[geomid, 11])
virtual_lin_mom = wp.vec3(
density * virtual_mass[0] * l_lin[0],
density * virtual_mass[1] * l_lin[1],
density * virtual_mass[2] * l_lin[2],
)
virtual_ang_mom = wp.vec3(
density * virtual_inertia[0] * l_ang[0],
density * virtual_inertia[1] * l_ang[1],
density * virtual_inertia[2] * l_ang[2],
)
added_mass_force = wp.cross(virtual_lin_mom, l_ang)
added_mass_torque = wp.cross(virtual_lin_mom, l_lin) + wp.cross(virtual_ang_mom, l_ang)
lfrc_force += added_mass_force
lfrc_torque += added_mass_torque
# lift force orthogonal to velocity from Kutta-Joukowski theorem
magnus_coef = geom_fluid[geomid, 5]
kutta_coef = geom_fluid[geomid, 4]
blunt_drag_coef = geom_fluid[geomid, 1]
slender_drag_coef = geom_fluid[geomid, 2]
ang_drag_coef = geom_fluid[geomid, 3]
volume = wp.static(4.0 / 3.0 * wp.pi) * semiaxes[0] * semiaxes[1] * semiaxes[2]
d_max = wp.max(wp.max(semiaxes[0], semiaxes[1]), semiaxes[2])
d_min = wp.min(wp.min(semiaxes[0], semiaxes[1]), semiaxes[2])
d_mid = semiaxes[0] + semiaxes[1] + semiaxes[2] - d_max - d_min
A_max = wp.pi * d_max * d_mid
lin_speed = wp.length(l_lin)
magnus_force = wp.cross(l_ang, l_lin) * (magnus_coef * density * volume)
s12 = semiaxes[1] * semiaxes[2]
s20 = semiaxes[2] * semiaxes[0]
s01 = semiaxes[0] * semiaxes[1]
proj_denom = _pow4(s12) * _pow2(l_lin[0]) + _pow4(s20) * _pow2(l_lin[1]) + _pow4(s01) * _pow2(l_lin[2])
proj_num = _pow2(s12 * l_lin[0]) + _pow2(s20 * l_lin[1]) + _pow2(s01 * l_lin[2])
A_proj = 0.0
cos_alpha = 0.0
if proj_num > MJ_MINVAL and proj_denom > MJ_MINVAL:
A_proj = wp.pi * wp.sqrt(proj_denom / wp.max(MJ_MINVAL, proj_num))
if lin_speed > MJ_MINVAL:
cos_alpha = proj_num / wp.max(MJ_MINVAL, lin_speed * proj_denom)
norm = wp.vec3(
_pow2(s12) * l_lin[0],
_pow2(s20) * l_lin[1],
_pow2(s01) * l_lin[2],
)
kutta_force = wp.vec3(0.0)
if density > 0.0 and kutta_coef != 0.0 and lin_speed > MJ_MINVAL:
kutta_circ = wp.cross(norm, l_lin) * (kutta_coef * density * cos_alpha * A_proj)
kutta_force = wp.cross(kutta_circ, l_lin)
eq_sphere_D = wp.static(2.0 / 3.0) * (semiaxes[0] + semiaxes[1] + semiaxes[2])
lin_visc_force_coef = wp.static(3.0 * wp.pi) * eq_sphere_D
lin_visc_torq_coef = wp.pi * eq_sphere_D * eq_sphere_D * eq_sphere_D
I_max = wp.static(8.0 / 15.0 * wp.pi) * d_mid * _pow4(d_max)
II0 = _ellipsoid_max_moment(semiaxes, 0)
II1 = _ellipsoid_max_moment(semiaxes, 1)
II2 = _ellipsoid_max_moment(semiaxes, 2)
mom_visc = wp.vec3(
l_ang[0] * (ang_drag_coef * II0 + slender_drag_coef * (I_max - II0)),
l_ang[1] * (ang_drag_coef * II1 + slender_drag_coef * (I_max - II1)),
l_ang[2] * (ang_drag_coef * II2 + slender_drag_coef * (I_max - II2)),
)
drag_lin_coef = viscosity * lin_visc_force_coef + density * lin_speed * (
A_proj * blunt_drag_coef + slender_drag_coef * (A_max - A_proj)
)
drag_ang_coef = viscosity * lin_visc_torq_coef + density * wp.length(mom_visc)
lfrc_torque -= drag_ang_coef * l_ang
lfrc_force += magnus_force + kutta_force - drag_lin_coef * l_lin
lfrc_torque *= coef
lfrc_force *= coef
# map force/torque from local to world frame: lfrc -> bfrc
torque_global += geom_rot @ lfrc_torque
force_global += geom_rot @ lfrc_force
fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global)
return
l_ang = rotT @ ang_global
l_lin = rotT @ lin_com
if wind[0] or wind[1] or wind[2]:
# subtract translational component from body velocity
lvel_force -= rotT @ wind
l_lin -= rotT @ wind
lfrc_torque = wp.vec3(0.0)
lfrc_force = wp.vec3(0.0)
@@ -284,54 +459,54 @@ def _box_fluid(
box2 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[1] - inertia[2]) * scl)
if has_viscosity:
# diameter of sphere approximation
diam = (box0 + box1 + box2) / 3.0
# angular viscosity
lfrc_torque = -lvel_torque * wp.pow(diam, 3.0) * wp.pi * viscosity
# linear viscosity
lfrc_force = -3.0 * lvel_force * diam * wp.pi * viscosity
lfrc_torque = -l_ang * wp.pow(diam, 3.0) * wp.pi * viscosity
lfrc_force = -3.0 * l_lin * diam * wp.pi * viscosity
if has_density:
# force
lfrc_force -= wp.vec3(
0.5 * density * box1 * box2 * wp.abs(lvel_force[0]) * lvel_force[0],
0.5 * density * box0 * box2 * wp.abs(lvel_force[1]) * lvel_force[1],
0.5 * density * box0 * box1 * wp.abs(lvel_force[2]) * lvel_force[2],
0.5 * density * box1 * box2 * wp.abs(l_lin[0]) * l_lin[0],
0.5 * density * box0 * box2 * wp.abs(l_lin[1]) * l_lin[1],
0.5 * density * box0 * box1 * wp.abs(l_lin[2]) * l_lin[2],
)
# torque
scl = density / 64.0
box0_pow4 = wp.pow(box0, 4.0)
box1_pow4 = wp.pow(box1, 4.0)
box2_pow4 = wp.pow(box2, 4.0)
lfrc_torque -= wp.vec3(
box0 * (box1_pow4 + box2_pow4) * wp.abs(lvel_torque[0]) * lvel_torque[0] * scl,
box1 * (box0_pow4 + box2_pow4) * wp.abs(lvel_torque[1]) * lvel_torque[1] * scl,
box2 * (box0_pow4 + box1_pow4) * wp.abs(lvel_torque[2]) * lvel_torque[2] * scl,
box0 * (box1_pow4 + box2_pow4) * wp.abs(l_ang[0]) * l_ang[0] * scl,
box1 * (box0_pow4 + box2_pow4) * wp.abs(l_ang[1]) * l_ang[1] * scl,
box2 * (box0_pow4 + box1_pow4) * wp.abs(l_ang[2]) * l_ang[2] * scl,
)
# rotate to global orientation: lfrc -> bfrc
torque = rot @ lfrc_torque
force = rot @ lfrc_force
torque_global = rot @ lfrc_torque
force_global = rot @ lfrc_force
fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force, torque)
fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global)
def _fluid(m: Model, d: Data):
wp.launch(
_box_fluid,
_fluid_force,
dim=(d.nworld, m.nbody),
inputs=[
m.opt.wind,
m.opt.density,
m.opt.viscosity,
m.body_rootid,
m.body_geomnum,
m.body_geomadr,
m.body_mass,
m.body_inertia,
m.body_fluid_ellipsoid,
m.geom_type,
m.geom_size,
m.geom_fluid,
d.xipos,
d.ximat,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
],
@@ -340,8 +515,6 @@ def _fluid(m: Model, d: Data):
],
)
# TODO(team): ellipsoid fluid model
support.apply_ft(m, d, d.fluid_applied, d.qfrc_fluid, False)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+19 -18
View File
@@ -1776,7 +1776,7 @@ def _sensor_acc(
sensor_adr_to_contact_adr: wp.array(dtype=int),
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
xpos_in: wp.array2d(dtype=wp.vec3),
xipos_in: wp.array2d(dtype=wp.vec3),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -1866,7 +1866,7 @@ def _sensor_acc(
contact_forcetorque = support.contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
nacon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
@@ -1896,7 +1896,7 @@ def _sensor_acc(
contact_forcetorque = support.contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
nacon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
@@ -1973,7 +1973,7 @@ def _sensor_acc(
contact_forcetorque = support.contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
nacon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
@@ -2082,7 +2082,7 @@ def _sensor_touch(
sensor_adr: wp.array(dtype=int),
sensor_touch_adr: wp.array(dtype=int),
# Data in:
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
contact_pos_in: wp.array(dtype=wp.vec3),
@@ -2097,7 +2097,7 @@ def _sensor_touch(
):
conid, sensortouchadrid = wp.tid()
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
sensorid = sensor_touch_adr[sensortouchadrid]
@@ -2182,7 +2182,7 @@ def _sensor_tactile(
taxel_vertadr: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int),
# Data in:
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
@@ -2194,7 +2194,7 @@ def _sensor_tactile(
):
conid, taxelid = wp.tid()
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
worldid = contact_worldid_in[conid]
@@ -2308,7 +2308,7 @@ def _contact_match(
sensor_contact_adr: wp.array(dtype=int),
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
contact_dist_in: wp.array(dtype=float),
@@ -2329,7 +2329,7 @@ def _contact_match(
contactsensorid, contactid = wp.tid()
sensorid = sensor_contact_adr[contactsensorid]
if contactid >= ncon_in[0]:
if contactid >= nacon_in[0]:
return
# sensor information
@@ -2403,7 +2403,7 @@ def _contact_match(
contact_force = support.contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
nacon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
@@ -2422,6 +2422,7 @@ def _contact_match(
sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = dir
@cache_kernel
def _contact_sort(maxmatch: int):
@nested_kernel(module="unique", enable_backward=False)
def contact_sort(
@@ -2466,7 +2467,7 @@ def sensor_acc(m: Model, d: Data):
wp.launch(
_sensor_touch,
dim=(d.nconmax, m.sensor_touch_adr.size),
dim=(d.naconmax, m.sensor_touch_adr.size),
inputs=[
m.opt.cone,
m.geom_bodyid,
@@ -2476,7 +2477,7 @@ def sensor_acc(m: Model, d: Data):
m.sensor_objid,
m.sensor_adr,
m.sensor_touch_adr,
d.ncon,
d.nacon,
d.site_xpos,
d.site_xmat,
d.contact.pos,
@@ -2494,7 +2495,7 @@ def sensor_acc(m: Model, d: Data):
wp.launch(
_sensor_tactile,
dim=(d.nconmax, m.nsensortaxel),
dim=(d.naconmax, m.nsensortaxel),
inputs=[
m.body_rootid,
m.body_weldid,
@@ -2518,7 +2519,7 @@ def sensor_acc(m: Model, d: Data):
m.geom_plugin_index,
m.taxel_vertadr,
m.taxel_sensorid,
d.ncon,
d.nacon,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
@@ -2539,7 +2540,7 @@ def sensor_acc(m: Model, d: Data):
wp.launch(
_contact_match,
dim=(m.sensor_contact_adr.size, d.nconmax),
dim=(m.sensor_contact_adr.size, d.naconmax),
inputs=[
m.opt.cone,
m.opt.contact_sensor_maxmatch,
@@ -2554,7 +2555,7 @@ def sensor_acc(m: Model, d: Data):
m.sensor_intprm,
m.sensor_contact_adr,
d.njmax,
d.ncon,
d.nacon,
d.site_xpos,
d.site_xmat,
d.contact.dist,
@@ -2616,7 +2617,7 @@ def sensor_acc(m: Model, d: Data):
m.sensor_acc_adr,
m.sensor_adr_to_contact_adr,
d.njmax,
d.ncon,
d.nacon,
d.xpos,
d.xipos,
d.geom_xpos,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+9 -13
View File
@@ -1243,7 +1243,7 @@ def _cfrc_ext_contact(
geom_bodyid: wp.array(dtype=int),
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
subtree_com_in: wp.array2d(dtype=wp.vec3),
contact_pos_in: wp.array(dtype=wp.vec3),
contact_frame_in: wp.array(dtype=wp.mat33),
@@ -1258,7 +1258,7 @@ def _cfrc_ext_contact(
):
contactid = wp.tid()
if contactid >= ncon_in[0]:
if contactid >= nacon_in[0]:
return
geom = contact_geom_in[contactid]
@@ -1274,7 +1274,7 @@ def _cfrc_ext_contact(
force = support.contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
nacon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
@@ -1338,13 +1338,13 @@ def rne_postconstraint(m: Model, d: Data):
# cfrc_ext += contacts
wp.launch(
_cfrc_ext_contact,
dim=(d.nconmax,),
dim=(d.naconmax,),
inputs=[
m.opt.cone,
m.body_rootid,
m.geom_bodyid,
d.njmax,
d.ncon,
d.nacon,
d.subtree_com,
d.contact.pos,
d.contact.frame,
@@ -2077,7 +2077,7 @@ def _transmission_body_moment(
actuator_trnid: wp.array(dtype=wp.vec2i),
actuator_trntype_body_adr: wp.array(dtype=int),
# Data in:
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector),
contact_dist_in: wp.array(dtype=float),
@@ -2097,7 +2097,7 @@ def _transmission_body_moment(
actid = actuator_trntype_body_adr[trnbodyid]
bodyid = actuator_trnid[actid][0]
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
worldid = contact_worldid_in[conid]
@@ -2247,11 +2247,7 @@ def transmission(m: Model, d: Data):
# compute moments
wp.launch(
_transmission_body_moment,
dim=(
m.actuator_trntype_body_adr.size,
d.nconmax,
m.nv,
),
dim=(m.actuator_trntype_body_adr.size, d.naconmax, m.nv),
inputs=[
m.opt.cone,
m.body_parentid,
@@ -2260,7 +2256,7 @@ def transmission(m: Model, d: Data):
m.geom_bodyid,
m.actuator_trnid,
m.actuator_trntype_body_adr,
d.ncon,
d.nacon,
d.subtree_com,
d.cdof,
d.contact.dist,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+16 -16
View File
@@ -501,7 +501,7 @@ def linesearch_parallel_fused(
opt_ls_parallel_min_step: float,
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int),
@@ -561,7 +561,7 @@ def linesearch_parallel_fused(
# extract contact info
conid = efc_id_in[worldid, efcid]
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
continue
efcid0 = contact_efc_address_in[conid, 0]
@@ -654,7 +654,7 @@ def _linesearch_parallel(m: types.Model, d: types.Data):
m.opt.impratio,
m.opt.ls_parallel_min_step,
d.njmax,
d.ncon,
d.nacon,
d.ne,
d.nf,
d.nefc,
@@ -771,7 +771,7 @@ def linesearch_prepare_quad(
# Model:
opt_impratio: wp.array(dtype=float),
# Data in:
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int),
contact_friction_in: wp.array(dtype=types.vec5),
contact_dim_in: wp.array(dtype=int),
@@ -805,7 +805,7 @@ def linesearch_prepare_quad(
# extract contact info
conid = efc_id_in[worldid, efcid]
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
efcid0 = contact_efc_address_in[conid, 0]
@@ -951,7 +951,7 @@ def _linesearch(m: types.Model, d: types.Data):
dim=(d.nworld, d.njmax),
inputs=[
m.opt.impratio,
d.ncon,
d.nacon,
d.nefc,
d.contact.friction,
d.contact.dim,
@@ -1064,7 +1064,7 @@ def update_constraint_efc(
# Model:
opt_impratio: wp.array(dtype=float),
# Data in:
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int),
@@ -1129,7 +1129,7 @@ def update_constraint_efc(
else: # elliptic friction cone contact
conid = efc_id_in[worldid, efcid]
if conid >= ncon_in[0]:
if conid >= nacon_in[0]:
return
dim = contact_dim_in[conid]
@@ -1264,7 +1264,7 @@ def _update_constraint(m: types.Model, d: types.Data):
dim=(d.nworld, d.njmax),
inputs=[
m.opt.impratio,
d.ncon,
d.nacon,
d.ne,
d.nf,
d.nefc,
@@ -1468,8 +1468,8 @@ def update_gradient_JTCJ(
dof_tri_row: wp.array(dtype=int),
dof_tri_col: wp.array(dtype=int),
# Data in:
nconmax_in: int,
ncon_in: wp.array(dtype=int),
naconmax_in: int,
nacon_in: wp.array(dtype=int),
contact_dist_in: wp.array(dtype=float),
contact_includemargin_in: wp.array(dtype=float),
contact_friction_in: wp.array(dtype=types.vec5),
@@ -1495,7 +1495,7 @@ def update_gradient_JTCJ(
for i in range(nblocks_perblock):
conid = conid_start + i * dim_block
if conid >= min(ncon_in[0], nconmax_in):
if conid >= min(nacon_in[0], naconmax_in):
return
worldid = contact_worldid_in[conid]
@@ -1727,9 +1727,9 @@ def _update_gradient(m: types.Model, d: types.Data):
dim_block = ceil((sm_count * 6 * 256) / m.dof_tri_row.size)
else:
# fall back for CPU
dim_block = d.nconmax
dim_block = d.naconmax
nblocks_perblock = int((d.nconmax + dim_block - 1) / dim_block)
nblocks_perblock = int((d.naconmax + dim_block - 1) / dim_block)
wp.launch(
update_gradient_JTCJ,
@@ -1738,8 +1738,8 @@ def _update_gradient(m: types.Model, d: types.Data):
m.opt.impratio,
m.dof_tri_row,
m.dof_tri_col,
d.nconmax,
d.ncon,
d.naconmax,
d.nacon,
d.contact.dist,
d.contact.includemargin,
d.contact.friction,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+6 -6
View File
@@ -296,7 +296,7 @@ def contact_force_fn(
opt_cone: int,
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int),
@@ -312,7 +312,7 @@ def contact_force_fn(
condim = contact_dim_in[contact_id]
efc_address = contact_efc_address_in[contact_id, 0]
if contact_id >= 0 and contact_id <= ncon_in[0] and efc_address >= 0:
if contact_id >= 0 and contact_id <= nacon_in[0] and efc_address >= 0:
if opt_cone == ConeType.PYRAMIDAL:
force = _decode_pyramid(
njmax_in,
@@ -341,7 +341,7 @@ def contact_force_kernel(
opt_cone: int,
# Data in:
njmax_in: int,
ncon_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int),
@@ -358,7 +358,7 @@ def contact_force_kernel(
contactid = contact_ids[tid]
if contactid >= ncon_in[0]:
if contactid >= nacon_in[0]:
return
worldid = contact_worldid_in[contactid]
@@ -366,7 +366,7 @@ def contact_force_kernel(
out[tid] = contact_force_fn(
opt_cone,
njmax_in,
ncon_in,
nacon_in,
contact_frame_in,
contact_friction_in,
contact_dim_in,
@@ -401,7 +401,7 @@ def contact_force(
inputs=[
m.opt.cone,
d.njmax,
d.ncon,
d.nacon,
d.contact.frame,
d.contact.friction,
d.contact.dim,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+44 -40
View File
@@ -807,6 +807,7 @@ class Model:
body_contype: OR over all geom contypes (nbody,)
body_conaffinity: OR over all geom conaffinities (nbody,)
body_gravcomp: antigravity force, units of body weight (nworld, nbody)
body_fluid_ellipsoid: does body use ellipsoid fluid (nbody,)
jnt_type: type of joint (JointType) (njnt,)
jnt_qposadr: start addr in 'qpos' for joint's data (njnt,)
jnt_dofadr: start addr in 'qvel' for joint's data (njnt,)
@@ -849,6 +850,7 @@ class Model:
geom_solref: constraint solver reference: contact (nworld, ngeom, mjNREF)
geom_solimp: constraint solver impedance: contact (nworld, ngeom, mjNIMP)
geom_size: geom-specific size parameters (ngeom, 3)
geom_fluid: fluid interaction parameters (ngeom, mjNFLUID)
geom_aabb: bounding box, (center, size) (ngeom, 6)
geom_rbound: radius of bounding sphere (nworld, ngeom,)
geom_pos: local position offset rel. to body (nworld, ngeom, 3)
@@ -1116,6 +1118,7 @@ class Model:
body_contype: wp.array(dtype=int)
body_conaffinity: wp.array(dtype=int)
body_gravcomp: wp.array2d(dtype=float)
body_fluid_ellipsoid: wp.array(dtype=bool)
jnt_type: wp.array(dtype=int)
jnt_qposadr: wp.array(dtype=int)
jnt_dofadr: wp.array(dtype=int)
@@ -1158,6 +1161,7 @@ class Model:
geom_solref: wp.array2d(dtype=wp.vec2)
geom_solimp: wp.array2d(dtype=vec5)
geom_size: wp.array2d(dtype=wp.vec3)
geom_fluid: wp.array2d(dtype=float)
geom_aabb: wp.array2d(dtype=wp.vec3)
geom_rbound: wp.array2d(dtype=float)
geom_pos: wp.array2d(dtype=wp.vec3)
@@ -1365,18 +1369,18 @@ class Contact:
"""Contact data.
Attributes:
dist: distance between nearest points; neg: penetration
pos: position of contact point: midpoint between geoms
frame: normal is in [0-2], points from geom[0] to geom[1]
includemargin: include if dist<includemargin=margin-gap
friction: tangent1, 2, spin, roll1, 2
solref: constraint solver reference, normal direction
solreffriction: constraint solver reference, friction directions
solimp: constraint solver impedance
dim: contact space dimensionality: 1, 3, 4 or 6
geom: geom ids; -1 for flex
efc_address: address in efc; -1: not included
worldid: world id
dist: distance between nearest points; neg: penetration (naconmax,)
pos: position of contact point: midpoint between geoms (naconmax, 3)
frame: normal is in [0-2], points from geom[0] to geom[1] (naconmax, 3, 3)
includemargin: include if dist<includemargin=margin-gap (naconmax,)
friction: tangent1, 2, spin, roll1, 2 (naconmax, 5)
solref: constraint solver reference, normal direction (naconmax, 2)
solreffriction: constraint solver reference, friction directions (naconmax, 2)
solimp: constraint solver impedance (naconmax, 5)
dim: contact space dimensionality: 1, 3, 4 or 6 (naconmax,)
geom: geom ids; -1 for flex (naconmax, 2)
efc_address: address in efc; -1: not included (naconmax, ncondim)
worldid: world id (naconmax,)
"""
dist: wp.array(dtype=float)
@@ -1399,10 +1403,10 @@ class Data:
Attributes:
nworld: number of worlds
nconmax: maximum number of contacts
naconmax: maximum number of contacts (shared across all worlds)
njmax: maximum number of constraints per world
solver_niter: number of solver iterations (nworld,)
ncon: number of detected contacts
nacon: number of detected contacts (across all worlds)
ne: number of equality constraints (nworld,)
ne_connect: number of equality connect constraints (nworld,)
ne_weld: number of equality weld constraints (nworld,)
@@ -1498,31 +1502,31 @@ class Data:
sap_cumulative_sum: broadphase context (nworld, ngeom)
sap_segment_index: broadphase context (requires nworld + 1) (nworld, 2)
dyn_geom_aabb: dynamic geometry axis-aligned bounding boxes (nworld, ngeom, 2)
collision_pair: collision pairs from broadphase (nconmax,)
collision_worldid: collision world ids from broadphase (nconmax,)
collision_pair: collision pairs from broadphase (naconmax,)
collision_worldid: collision world ids from broadphase (naconmax,)
ncollision: collision count from broadphase
epa_vert: vertices in EPA polytope in Minkowski space (nconmax, 5 + CCDiter)
epa_vert1: vertices in EPA polytope in geom 1 space (nconmax, 5 + CCDiter)
epa_vert2: vertices in EPA polytope in geom 2 space (nconmax, 5 + CCDiter)
epa_vert_index1: vertex indices in EPA polytope for geom 1 (nconmax, 5 + CCDiter)
epa_vert_index2: vertex indices in EPA polytope for geom 2 (nconmax, 5 + CCDiter)
epa_face: faces of polytope represented by three indices (nconmax, 6 + 5 * CCDiter)
epa_pr: projection of origin on polytope faces (nconmax, 6 + 5 * CCDiter)
epa_norm2: epa_pr * epa_pr (nconmax, 6 + 5 * CCDiter)
epa_index: index of face in polytope map (nconmax, 6 + 5 * CCDiter)
epa_map: status of faces in polytope (nconmax, 6 + 5 * CCDiter)
epa_horizon: index pair (i j) of edges on horizon (nconmax, 2 * 12)
multiccd_polygon: clipped contact surface (nconmax, 2 * max_npolygon)
multiccd_clipped: clipped contact surface (intermediate) (nconmax, 2 * max_npolygon)
multiccd_pnormal: plane normal of clipping polygon (nconmax, max_npolygon)
multiccd_pdist: plane distance of clipping polygon (nconmax, max_npolygon)
multiccd_idx1: list of normal index candidates for Geom 1 (nconmax, max_meshdegree)
multiccd_idx2: list of normal index candidates for Geom 2 (nconmax, max_meshdegree)
multiccd_n1: list of normal candidates for Geom 1 (nconmax, max_meshdegree)
multiccd_n2: list of normal candidates for Geom 1 (nconmax, max_meshdegree)
multiccd_endvert: list of edge vertices candidates (nconmax, max_meshdegree)
multiccd_face1: contact face (nconmax, max_npolygon)
multiccd_face2: contact face (nconmax, max_npolygon)
epa_vert: vertices in EPA polytope in Minkowski space (naconmax, 5 + CCDiter)
epa_vert1: vertices in EPA polytope in geom 1 space (naconmax, 5 + CCDiter)
epa_vert2: vertices in EPA polytope in geom 2 space (naconmax, 5 + CCDiter)
epa_vert_index1: vertex indices in EPA polytope for geom 1 (naconmax, 5 + CCDiter)
epa_vert_index2: vertex indices in EPA polytope for geom 2 (naconmax, 5 + CCDiter)
epa_face: faces of polytope represented by three indices (naconmax, 6 + 5 * CCDiter)
epa_pr: projection of origin on polytope faces (naconmax, 6 + 5 * CCDiter)
epa_norm2: epa_pr * epa_pr (naconmax, 6 + 5 * CCDiter)
epa_index: index of face in polytope map (naconmax, 6 + 5 * CCDiter)
epa_map: status of faces in polytope (naconmax, 6 + 5 * CCDiter)
epa_horizon: index pair (i j) of edges on horizon (naconmax, 2 * 12)
multiccd_polygon: clipped contact surface (naconmax, 2 * max_npolygon)
multiccd_clipped: clipped contact surface (intermediate) (naconmax, 2 * max_npolygon)
multiccd_pnormal: plane normal of clipping polygon (naconmax, max_npolygon)
multiccd_pdist: plane distance of clipping polygon (naconmax, max_npolygon)
multiccd_idx1: list of normal index candidates for Geom 1 (naconmax, max_meshdegree)
multiccd_idx2: list of normal index candidates for Geom 2 (naconmax, max_meshdegree)
multiccd_n1: list of normal candidates for Geom 1 (naconmax, max_meshdegree)
multiccd_n2: list of normal candidates for Geom 1 (naconmax, max_meshdegree)
multiccd_endvert: list of edge vertices candidates (naconmax, max_meshdegree)
multiccd_face1: contact face (naconmax, max_npolygon)
multiccd_face2: contact face (naconmax, max_npolygon)
cacc: com-based acceleration (nworld, nbody, 6)
cfrc_int: com-based interaction force with parent (nworld, nbody, 6)
cfrc_ext: com-based external force on body (nworld, nbody, 6)
@@ -1554,10 +1558,10 @@ class Data:
"""
nworld: int # warp only
nconmax: int # warp only
naconmax: int # warp only
njmax: int # warp only
solver_niter: wp.array(dtype=int)
ncon: wp.array(dtype=int)
nacon: wp.array(dtype=int)
ne: wp.array(dtype=int)
ne_connect: wp.array(dtype=int) # warp only
ne_weld: wp.array(dtype=int) # warp only
mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
Vendored
+35 -9
View File
@@ -157,6 +157,13 @@ def kernel(
tid = wp.tid()
b[tid] = a[tid] + 1.0
@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.
@@ -168,15 +175,27 @@ def kernel(
Returns:
The registered kernel.
"""
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 = f.__qualname__
parts = [part for part in qualname.split(".") if part != "<locals>"]
outer_functions = parts[:-1]
module = get_module(".".join([f.__module__] + outer_functions))
return wp.kernel(f, enable_backward=enable_backward, module=module)
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 = 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 = {}
@@ -186,7 +205,14 @@ def cache_kernel(func):
# caching kernels to avoid crashes in graph_conditional code
@functools.wraps(func)
def wrapper(*args):
key = tuple(a.size if hasattr(a, "size") else hash(a) for a in args) + (hash(func.__name__),)
def _hash_arg(a):
if hasattr(a, "size"):
return a.size
if isinstance(a, list):
return hash(tuple(a))
return hash(a)
key = tuple(_hash_arg(a) for a in args) + (hash(func.__name__),)
if key not in _KERNEL_CACHE:
_KERNEL_CACHE[key] = func(*args)
return _KERNEL_CACHE[key]
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+1 -1
View File
@@ -71,7 +71,7 @@ def key_callback(key: int) -> None:
def _load_model(path: epath.Path) -> mujoco.MjModel:
if not path.exists():
resource_path = epath.resource_path("mujoco_warp") / path
resource_path = epath.resource_path("mjx") / "third_party/mujoco_warp" / path
if not resource_path.exists():
raise FileNotFoundError(f"file not found: {path}\nalso tried: {resource_path}")
path = resource_path
mjx/mujoco/mjx/warp/collision_driver.py
+10 -12
View File
@@ -112,7 +112,7 @@ def _collision_shim(
opt__sdf_initpoints: int,
opt__sdf_iterations: int,
# Data
nconmax: int,
naconmax: int,
collision_pair: wp.array(dtype=wp.vec2i),
collision_pairid: wp.array(dtype=int),
collision_worldid: wp.array(dtype=int),
@@ -140,8 +140,8 @@ def _collision_shim(
multiccd_pdist: wp.array2d(dtype=float),
multiccd_pnormal: wp.array2d(dtype=wp.vec3),
multiccd_polygon: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int),
ncon: wp.array(dtype=int),
sap_cumulative_sum: wp.array2d(dtype=int),
sap_projection_lower: wp.array3d(dtype=float),
sap_projection_upper: wp.array2d(dtype=float),
@@ -266,9 +266,9 @@ def _collision_shim(
_d.multiccd_pdist = multiccd_pdist
_d.multiccd_pnormal = multiccd_pnormal
_d.multiccd_polygon = multiccd_polygon
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision
_d.ncon = ncon
_d.nconmax = nconmax
_d.sap_cumulative_sum = sap_cumulative_sum
_d.sap_projection_lower = sap_projection_lower
_d.sap_projection_upper = sap_projection_upper
@@ -308,8 +308,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'multiccd_pdist': d._impl.multiccd_pdist.shape,
'multiccd_pnormal': d._impl.multiccd_pnormal.shape,
'multiccd_polygon': d._impl.multiccd_polygon.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape,
'ncon': d._impl.ncon.shape,
'sap_cumulative_sum': d._impl.sap_cumulative_sum.shape,
'sap_projection_lower': d._impl.sap_projection_lower.shape,
'sap_projection_upper': d._impl.sap_projection_upper.shape,
@@ -361,8 +361,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'multiccd_pdist',
'multiccd_pnormal',
'multiccd_polygon',
'nacon',
'ncollision',
'ncon',
'sap_cumulative_sum',
'sap_projection_lower',
'sap_projection_upper',
@@ -448,7 +448,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.opt._impl.legacy_gjk,
m.opt._impl.sdf_initpoints,
m.opt._impl.sdf_iterations,
d._impl.nconmax,
d._impl.naconmax,
d._impl.collision_pair,
d._impl.collision_pairid,
d._impl.collision_worldid,
@@ -476,8 +476,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
d._impl.multiccd_pdist,
d._impl.multiccd_pnormal,
d._impl.multiccd_polygon,
d._impl.nacon,
d._impl.ncollision,
d._impl.ncon,
d._impl.sap_cumulative_sum,
d._impl.sap_projection_lower,
d._impl.sap_projection_upper,
@@ -524,8 +524,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'_impl.multiccd_pdist': out[24],
'_impl.multiccd_pnormal': out[25],
'_impl.multiccd_polygon': out[26],
'_impl.ncollision': out[27],
'_impl.ncon': out[28],
'_impl.nacon': out[27],
'_impl.ncollision': out[28],
'_impl.sap_cumulative_sum': out[29],
'_impl.sap_projection_lower': out[30],
'_impl.sap_projection_upper': out[31],
@@ -551,8 +551,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable
def collision(m: types.Model, d: types.Data):
return _collision_jax_impl(m, d)
@collision.def_vmap
@ffi.marshal_custom_vmap
def collision_vmap(unused_axis_size, is_batched, m, d):
mjx/mujoco/mjx/warp/ffi.py
+1 -1
View File
@@ -360,7 +360,7 @@ def marshal_custom_vmap(vmap_func):
# Check leading dimensions.
jax.tree.map_with_path(
lambda path, x: _check_leading_dim(
path, x, d_broadcast.qpos.shape[0], d._impl.nconmax, d._impl.njmax # pylint: disable=protected-access
path, x, d_broadcast.qpos.shape[0], d._impl.naconmax, d._impl.njmax # pylint: disable=protected-access
),
d_broadcast,
)
mjx/mujoco/mjx/warp/forward.py
+32 -24
View File
@@ -74,6 +74,7 @@ def _forward_shim(
block_dim: mjwp_types.BlockDim,
body_dofadr: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int),
body_fluid_ellipsoid: wp.array(dtype=bool),
body_geomadr: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int),
body_gravcomp: wp.array2d(dtype=float),
@@ -144,6 +145,7 @@ def _forward_shim(
geom_bodyid: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_fluid: wp.array2d(dtype=float),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_gap: wp.array2d(dtype=float),
geom_group: wp.array(dtype=int),
@@ -351,7 +353,7 @@ def _forward_shim(
opt__wind: wp.array(dtype=wp.vec3),
stat__meaninertia: float,
# Data
nconmax: int,
naconmax: int,
njmax: int,
act: wp.array2d(dtype=float),
act_dot: wp.array2d(dtype=float),
@@ -409,8 +411,8 @@ def _forward_shim(
multiccd_pdist: wp.array2d(dtype=float),
multiccd_pnormal: wp.array2d(dtype=wp.vec3),
multiccd_polygon: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int),
ncon: wp.array(dtype=int),
ne: wp.array(dtype=int),
ne_connect: wp.array(dtype=int),
ne_jnt: wp.array(dtype=int),
@@ -558,6 +560,7 @@ def _forward_shim(
_m.block_dim = block_dim
_m.body_dofadr = body_dofadr
_m.body_dofnum = body_dofnum
_m.body_fluid_ellipsoid = body_fluid_ellipsoid
_m.body_geomadr = body_geomadr
_m.body_geomnum = body_geomnum
_m.body_gravcomp = body_gravcomp
@@ -628,6 +631,7 @@ def _forward_shim(
_m.geom_bodyid = geom_bodyid
_m.geom_condim = geom_condim
_m.geom_dataid = geom_dataid
_m.geom_fluid = geom_fluid
_m.geom_friction = geom_friction
_m.geom_gap = geom_gap
_m.geom_group = geom_group
@@ -936,9 +940,9 @@ def _forward_shim(
_d.multiccd_pdist = multiccd_pdist
_d.multiccd_pnormal = multiccd_pnormal
_d.multiccd_polygon = multiccd_polygon
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision
_d.ncon = ncon
_d.nconmax = nconmax
_d.ne = ne
_d.ne_connect = ne_connect
_d.ne_jnt = ne_jnt
@@ -1071,8 +1075,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'multiccd_pdist': d._impl.multiccd_pdist.shape,
'multiccd_pnormal': d._impl.multiccd_pnormal.shape,
'multiccd_polygon': d._impl.multiccd_polygon.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape,
'ncon': d._impl.ncon.shape,
'ne': d._impl.ne.shape,
'ne_connect': d._impl.ne_connect.shape,
'ne_jnt': d._impl.ne_jnt.shape,
@@ -1251,8 +1255,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'multiccd_pdist',
'multiccd_pnormal',
'multiccd_polygon',
'nacon',
'ncollision',
'ncon',
'ne',
'ne_connect',
'ne_jnt',
@@ -1399,6 +1403,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.block_dim,
m.body_dofadr,
m.body_dofnum,
m._impl.body_fluid_ellipsoid,
m.body_geomadr,
m.body_geomnum,
m.body_gravcomp,
@@ -1469,6 +1474,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.geom_bodyid,
m.geom_condim,
m.geom_dataid,
m.geom_fluid,
m.geom_friction,
m.geom_gap,
m.geom_group,
@@ -1675,7 +1681,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.opt.viscosity,
m.opt.wind,
m.stat.meaninertia,
d._impl.nconmax,
d._impl.naconmax,
d._impl.njmax,
d.act,
d.act_dot,
@@ -1733,8 +1739,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.multiccd_pdist,
d._impl.multiccd_pnormal,
d._impl.multiccd_polygon,
d._impl.nacon,
d._impl.ncollision,
d._impl.ncon,
d._impl.ne,
d._impl.ne_connect,
d._impl.ne_jnt,
@@ -1908,8 +1914,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'_impl.multiccd_pdist': out[53],
'_impl.multiccd_pnormal': out[54],
'_impl.multiccd_polygon': out[55],
'_impl.ncollision': out[56],
'_impl.ncon': out[57],
'_impl.nacon': out[56],
'_impl.ncollision': out[57],
'_impl.ne': out[58],
'_impl.ne_connect': out[59],
'_impl.ne_jnt': out[60],
@@ -2033,8 +2039,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable
def forward(m: types.Model, d: types.Data):
return _forward_jax_impl(m, d)
@forward.def_vmap
@ffi.marshal_custom_vmap
def forward_vmap(unused_axis_size, is_batched, m, d):
@@ -2093,6 +2097,7 @@ def _step_shim(
block_dim: mjwp_types.BlockDim,
body_dofadr: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int),
body_fluid_ellipsoid: wp.array(dtype=bool),
body_geomadr: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int),
body_gravcomp: wp.array2d(dtype=float),
@@ -2163,6 +2168,7 @@ def _step_shim(
geom_bodyid: wp.array(dtype=int),
geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_fluid: wp.array2d(dtype=float),
geom_friction: wp.array2d(dtype=wp.vec3),
geom_gap: wp.array2d(dtype=float),
geom_group: wp.array(dtype=int),
@@ -2371,7 +2377,7 @@ def _step_shim(
opt__wind: wp.array(dtype=wp.vec3),
stat__meaninertia: float,
# Data
nconmax: int,
naconmax: int,
njmax: int,
act: wp.array2d(dtype=float),
act_dot: wp.array2d(dtype=float),
@@ -2432,8 +2438,8 @@ def _step_shim(
multiccd_pdist: wp.array2d(dtype=float),
multiccd_pnormal: wp.array2d(dtype=wp.vec3),
multiccd_polygon: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int),
ncon: wp.array(dtype=int),
ne: wp.array(dtype=int),
ne_connect: wp.array(dtype=int),
ne_jnt: wp.array(dtype=int),
@@ -2590,6 +2596,7 @@ def _step_shim(
_m.block_dim = block_dim
_m.body_dofadr = body_dofadr
_m.body_dofnum = body_dofnum
_m.body_fluid_ellipsoid = body_fluid_ellipsoid
_m.body_geomadr = body_geomadr
_m.body_geomnum = body_geomnum
_m.body_gravcomp = body_gravcomp
@@ -2660,6 +2667,7 @@ def _step_shim(
_m.geom_bodyid = geom_bodyid
_m.geom_condim = geom_condim
_m.geom_dataid = geom_dataid
_m.geom_fluid = geom_fluid
_m.geom_friction = geom_friction
_m.geom_gap = geom_gap
_m.geom_group = geom_group
@@ -2972,9 +2980,9 @@ def _step_shim(
_d.multiccd_pdist = multiccd_pdist
_d.multiccd_pnormal = multiccd_pnormal
_d.multiccd_polygon = multiccd_polygon
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision
_d.ncon = ncon
_d.nconmax = nconmax
_d.ne = ne
_d.ne_connect = ne_connect
_d.ne_jnt = ne_jnt
@@ -3119,8 +3127,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'multiccd_pdist': d._impl.multiccd_pdist.shape,
'multiccd_pnormal': d._impl.multiccd_pnormal.shape,
'multiccd_polygon': d._impl.multiccd_polygon.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape,
'ncon': d._impl.ncon.shape,
'ne': d._impl.ne.shape,
'ne_connect': d._impl.ne_connect.shape,
'ne_jnt': d._impl.ne_jnt.shape,
@@ -3311,8 +3319,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'multiccd_pdist',
'multiccd_pnormal',
'multiccd_polygon',
'nacon',
'ncollision',
'ncon',
'ne',
'ne_connect',
'ne_jnt',
@@ -3468,6 +3476,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.block_dim,
m.body_dofadr,
m.body_dofnum,
m._impl.body_fluid_ellipsoid,
m.body_geomadr,
m.body_geomnum,
m.body_gravcomp,
@@ -3538,6 +3547,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.geom_bodyid,
m.geom_condim,
m.geom_dataid,
m.geom_fluid,
m.geom_friction,
m.geom_gap,
m.geom_group,
@@ -3745,7 +3755,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.opt.viscosity,
m.opt.wind,
m.stat.meaninertia,
d._impl.nconmax,
d._impl.naconmax,
d._impl.njmax,
d.act,
d.act_dot,
@@ -3806,8 +3816,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.multiccd_pdist,
d._impl.multiccd_pnormal,
d._impl.multiccd_polygon,
d._impl.nacon,
d._impl.ncollision,
d._impl.ncon,
d._impl.ne,
d._impl.ne_connect,
d._impl.ne_jnt,
@@ -3993,8 +4003,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'_impl.multiccd_pdist': out[56],
'_impl.multiccd_pnormal': out[57],
'_impl.multiccd_polygon': out[58],
'_impl.ncollision': out[59],
'_impl.ncon': out[60],
'_impl.nacon': out[59],
'_impl.ncollision': out[60],
'_impl.ne': out[61],
'_impl.ne_connect': out[62],
'_impl.ne_jnt': out[63],
@@ -4127,8 +4137,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable
def step(m: types.Model, d: types.Data):
return _step_jax_impl(m, d)
@step.def_vmap
@ffi.marshal_custom_vmap
def step_vmap(unused_axis_size, is_batched, m, d):
mjx/mujoco/mjx/warp/smooth.py
-6
View File
@@ -42,7 +42,6 @@ _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
@@ -281,8 +280,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable
def kinematics(m: types.Model, d: types.Data):
return _kinematics_jax_impl(m, d)
@kinematics.def_vmap
@ffi.marshal_custom_vmap
def kinematics_vmap(unused_axis_size, is_batched, m, d):
@@ -309,7 +306,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
)
@ffi.format_args_for_warp
def _tendon_shim(
# Model
@@ -491,8 +487,6 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable
def tendon(m: types.Model, d: types.Data):
return _tendon_jax_impl(m, d)
@tendon.def_vmap
@ffi.marshal_custom_vmap
def tendon_vmap(unused_axis_size, is_batched, m, d):
mjx/mujoco/mjx/warp/test_util.py
+1 -1
View File
@@ -103,7 +103,7 @@ def make_data(
def _mjx_contact(dx, worldid: int):
keys = []
for i in range(dx._impl.ncon[0]):
for i in range(dx._impl.nacon[0]):
if dx._impl.contact__worldid[i] != worldid:
continue
mjx/mujoco/mjx/warp/types.py
+13 -23
View File
@@ -13,7 +13,6 @@
# limitations under the License.
# ==============================================================================
"""MJX Warp types.
DO NOT EDIT. This file is auto-generated.
"""
import dataclasses
@@ -23,10 +22,8 @@ from jax import tree_util
from jax.interpreters import batching
from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
PyTreeNode = mjx_dataclasses.PyTreeNode
@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class TileSet:
@@ -38,7 +35,6 @@ class TileSet:
adr: address of each tile in the set
size: size of all the tiles in this set
"""
adr: np.ndarray
size: int
@@ -59,7 +55,6 @@ class BlockDim:
TODO(team): experimental and may be removed
"""
actuator_velocity: int
cholesky_factorize: int
cholesky_factorize_solve: int
@@ -87,13 +82,10 @@ class BlockDim:
class StatisticWarp(PyTreeNode):
"""Derived fields from Statistic."""
meaninertia: float
class OptionWarp(PyTreeNode):
"""Derived fields from Option."""
broadphase: int
broadphase_filter: int
ccd_iterations: int
@@ -109,10 +101,8 @@ class OptionWarp(PyTreeNode):
sdf_initpoints: int
sdf_iterations: int
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
M_colind: np.ndarray
M_rowadr: np.ndarray
M_rownnz: np.ndarray
@@ -120,6 +110,7 @@ class ModelWarp(PyTreeNode):
actuator_moment_tiles_nv: Tuple[TileSet, ...]
actuator_trntype_body_adr: np.ndarray
block_dim: BlockDim
body_fluid_ellipsoid: np.ndarray
body_tree: Tuple[np.ndarray, ...]
collision_sensor_adr: np.ndarray
condim_max: int
@@ -218,10 +209,8 @@ class ModelWarp(PyTreeNode):
wrap_site_adr: np.ndarray
wrap_site_pair_adr: np.ndarray
class DataWarp(PyTreeNode):
"""Derived fields from Data."""
act_dot_rk: jax.Array
act_t0: jax.Array
act_vel_integration: jax.Array
@@ -317,9 +306,9 @@ class DataWarp(PyTreeNode):
multiccd_pdist: jax.Array
multiccd_pnormal: jax.Array
multiccd_polygon: jax.Array
nacon: jax.Array
naconmax: int
ncollision: jax.Array
ncon: jax.Array
nconmax: int
ne: jax.Array
ne_connect: jax.Array
ne_jnt: jax.Array
@@ -378,8 +367,6 @@ class DataWarp(PyTreeNode):
wrap_obj: jax.Array
wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
'collision_pair',
'collision_pairid',
@@ -419,16 +406,15 @@ DATA_NON_VMAP = {
'multiccd_pdist',
'multiccd_pnormal',
'multiccd_polygon',
'nacon',
'naconmax',
'ncollision',
'ncon',
'nconmax',
'njmax',
'nsolving',
'nworld',
'ray_bodyexclude',
}
def _to_elt(cont, _, d, axis):
return DataWarp(**{
f.name: (
@@ -562,9 +548,9 @@ _NDIM = {
'multiccd_pdist': 2,
'multiccd_pnormal': 3,
'multiccd_polygon': 3,
'nacon': 1,
'naconmax': 0,
'ncollision': 1,
'ncon': 1,
'nconmax': 0,
'ne': 1,
'ne_connect': 1,
'ne_jnt': 1,
@@ -697,6 +683,7 @@ _NDIM = {
'body_contype': 1,
'body_dofadr': 1,
'body_dofnum': 1,
'body_fluid_ellipsoid': 1,
'body_geomadr': 1,
'body_geomnum': 1,
'body_gravcomp': 2,
@@ -774,6 +761,7 @@ _NDIM = {
'geom_condim': 1,
'geom_contype': 1,
'geom_dataid': 1,
'geom_fluid': 2,
'geom_friction': 3,
'geom_gap': 2,
'geom_group': 1,
@@ -1155,9 +1143,9 @@ _BATCH_DIM = {
'multiccd_pdist': False,
'multiccd_pnormal': False,
'multiccd_polygon': False,
'nacon': False,
'naconmax': False,
'ncollision': False,
'ncon': False,
'nconmax': False,
'ne': True,
'ne_connect': True,
'ne_jnt': True,
@@ -1290,6 +1278,7 @@ _BATCH_DIM = {
'body_contype': False,
'body_dofadr': False,
'body_dofnum': False,
'body_fluid_ellipsoid': False,
'body_geomadr': False,
'body_geomnum': False,
'body_gravcomp': True,
@@ -1367,6 +1356,7 @@ _BATCH_DIM = {
'geom_condim': False,
'geom_contype': False,
'geom_dataid': False,
'geom_fluid': False,
'geom_friction': True,
'geom_gap': True,
'geom_group': False,