chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 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 :class: attention
- The attribute ``mjOption.apirate`` was unused and has been removed. - 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 MJX
^^^ ^^^
- Fix :github:issue:`2508`, ``qLD`` shapes mismatched mjModel during ``get_data_into``. - 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) 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( def _make_data_warp(
m: Union[types.Model, mujoco.MjModel], m: Union[types.Model, mujoco.MjModel],
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
nconmax: int = -1, nconmax: Optional[int] = None,
njmax: int = -1, njmax: Optional[int] = None,
) -> types.Data: ) -> types.Data:
"""Allocate and initialize Data for the Warp implementation.""" """Allocate and initialize Data for the Warp implementation."""
if not isinstance(m, mujoco.MjModel): if not isinstance(m, mujoco.MjModel):
@@ -841,7 +841,7 @@ def _make_data_warp(
raise RuntimeError('Warp is not installed.') raise RuntimeError('Warp is not installed.')
with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable 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) fields = _make_data_public_fields(m)
for k in fields: for k in fields:
@@ -876,7 +876,7 @@ def _make_data_warp(
# TODO(robotics-simulation): remove this warmup compilation once warp # TODO(robotics-simulation): remove this warmup compilation once warp
# stops unloading modules during XLA graph capture for tile kernels. # stops unloading modules during XLA graph capture for tile kernels.
# pylint: disable=undefined-variable # 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) mw = mjwp.put_model(m)
_ = mjwp.step(mw, dw) _ = mjwp.step(mw, dw)
# pylint: enable=undefined-variable # pylint: enable=undefined-variable
@@ -890,8 +890,8 @@ def make_data(
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
impl: Optional[Union[str, types.Impl]] = None, impl: Optional[Union[str, types.Impl]] = None,
_full_compat: bool = False, # pylint: disable=invalid-name _full_compat: bool = False, # pylint: disable=invalid-name
nconmax: int = -1, nconmax: Optional[int] = None,
njmax: int = -1, njmax: Optional[int] = None,
) -> types.Data: ) -> types.Data:
"""Allocate and initialize 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 _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 an experimental feature. Avoid using it for now. If using this flag, also
use _full_compat for put_model. use _full_compat for put_model.
nconmax: maximum number of contacts to allocate for warp nconmax: maximum number of contacts to allocate for warp across all worlds
njmax: maximum number of constraints to allocate for warp 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: Returns:
an initialized mjx.Data placed on device 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,)) self.assertEqual(dx[0].qpos.shape, (m.nq,))
if impl == 'warp': if impl == 'warp':
self.assertEqual(dx._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.nconmax,)) self.assertEqual(dx[0]._impl.contact__dist.shape, (dx._impl.naconmax,))
class FullCompatTest(parameterized.TestCase): 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 = {} trace = {}
ncon, nefc, solver_niter = [], [], [] nacon, nefc, solver_niter = [], [], []
center = wp.array([], dtype=wp.float32) center = wp.array([], dtype=wp.float32)
with warp_util.EventTracer(enabled=event_trace) as tracer: with warp_util.EventTracer(enabled=event_trace) as tracer:
@@ -142,7 +142,7 @@ def benchmark(
else: else:
trace = tracer.trace() trace = tracer.trace()
if measure_alloc: 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())) nefc.append(np.max(d.nefc.numpy()))
if measure_solver_niter: if measure_solver_niter:
solver_niter.append(d.solver_niter.numpy()) 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)) nsuccess = np.sum(~np.any(np.isnan(d.qpos.numpy()), axis=1))
run_duration = np.sum(time_vec) 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: 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), opt_ccd_tolerance: wp.array(dtype=float),
geom_type: wp.array(dtype=int), geom_type: wp.array(dtype=int),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
epa_vert_in: wp.array2d(dtype=wp.vec3), epa_vert_in: wp.array2d(dtype=wp.vec3),
epa_vert1_in: wp.array2d(dtype=wp.vec3), epa_vert1_in: wp.array2d(dtype=wp.vec3),
epa_vert2_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, x2: wp.vec3,
count: int, count: int,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -222,7 +222,7 @@ def ccd_kernel_builder(
for i in range(ncontact): for i in range(ncontact):
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
points[i], points[i],
frame, frame,
@@ -235,7 +235,7 @@ def ccd_kernel_builder(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -298,7 +298,7 @@ def ccd_kernel_builder(
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d(dtype=vec5),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_in: wp.array(dtype=wp.vec2i), 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_face1_in: wp.array2d(dtype=wp.vec3),
multiccd_face2_in: wp.array2d(dtype=wp.vec3), multiccd_face2_in: wp.array2d(dtype=wp.vec3),
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -503,7 +503,7 @@ def ccd_kernel_builder(
ncontact = eval_ccd_write_contact( ncontact = eval_ccd_write_contact(
opt_ccd_tolerance, opt_ccd_tolerance,
geom_type, geom_type,
nconmax_in, naconmax_in,
epa_vert_in, epa_vert_in,
epa_vert1_in, epa_vert1_in,
epa_vert2_in, epa_vert2_in,
@@ -541,7 +541,7 @@ def ccd_kernel_builder(
x1, x1,
geom2.pos, geom2.pos,
count, count,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -561,7 +561,7 @@ def ccd_kernel_builder(
eval_ccd_write_contact( eval_ccd_write_contact(
opt_ccd_tolerance, opt_ccd_tolerance,
geom_type, geom_type,
nconmax_in, naconmax_in,
epa_vert_in, epa_vert_in,
epa_vert1_in, epa_vert1_in,
epa_vert2_in, epa_vert2_in,
@@ -599,7 +599,7 @@ def ccd_kernel_builder(
geom1.pos, geom1.pos,
geom2.pos, geom2.pos,
0, 0,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_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)]: if m.geom_pair_type_count[upper_trid_index(len(GeomType), g1, g2)]:
wp.launch( wp.launch(
ccd_kernel_builder(m.opt.legacy_gjk, g1, g2, m.opt.ccd_iterations, True, 1e9, g1 == GeomType.HFIELD), 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=[ inputs=[
m.opt.ccd_tolerance, m.opt.ccd_tolerance,
m.geom_type, m.geom_type,
@@ -680,7 +680,7 @@ def convex_narrowphase(m: Model, d: Data):
m.pair_margin, m.pair_margin,
m.pair_gap, m.pair_gap,
m.pair_friction, m.pair_friction,
d.nconmax, d.naconmax,
d.geom_xpos, d.geom_xpos,
d.geom_xmat, d.geom_xmat,
d.collision_pair, d.collision_pair,
@@ -711,7 +711,7 @@ def convex_narrowphase(m: Model, d: Data):
d.multiccd_face2, d.multiccd_face2,
], ],
outputs=[ outputs=[
d.ncon, d.nacon,
d.contact.dist, d.contact.dist,
d.contact.pos, d.contact.pos,
d.contact.frame, 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 @wp.kernel
def _zero_ncon_ncollision( def _zero_nacon_ncollision(
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
ncollision_out: wp.array(dtype=int), ncollision_out: wp.array(dtype=int),
): ):
ncollision_out[0] = 0 ncollision_out[0] = 0
ncon_out[0] = 0 nacon_out[0] = 0
@wp.func @wp.func
@@ -292,7 +292,7 @@ def _add_geom_pair(
geom_type: wp.array(dtype=int), geom_type: wp.array(dtype=int),
nxn_pairid: wp.array(dtype=int), nxn_pairid: wp.array(dtype=int),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
geom1: int, geom1: int,
geom2: int, geom2: int,
@@ -306,7 +306,7 @@ def _add_geom_pair(
): ):
pairid = wp.atomic_add(ncollision_out, 0, 1) pairid = wp.atomic_add(ncollision_out, 0, 1)
if pairid >= nconmax_in: if pairid >= naconmax_in:
return return
type1 = geom_type[geom1] type1 = geom_type[geom1]
@@ -407,7 +407,7 @@ def _sap_broadphase(broadphase_filter):
nxn_pairid: wp.array(dtype=int), nxn_pairid: wp.array(dtype=int),
# Data in: # Data in:
nworld_in: int, nworld_in: int,
nconmax_in: int, naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore sap_sort_index_in: wp.array2d(dtype=int), # kernel_analyzer: ignore
@@ -455,7 +455,7 @@ def _sap_broadphase(broadphase_filter):
_add_geom_pair( _add_geom_pair(
geom_type, geom_type,
nxn_pairid, nxn_pairid,
nconmax_in, naconmax_in,
geom1, geom1,
geom2, geom2,
worldid, worldid,
@@ -581,7 +581,7 @@ def sap_broadphase(m: Model, d: Data):
m.geom_margin, m.geom_margin,
m.nxn_pairid, m.nxn_pairid,
d.nworld, d.nworld,
d.nconmax, d.naconmax,
d.geom_xpos, d.geom_xpos,
d.geom_xmat, d.geom_xmat,
d.sap_sort_index.reshape((-1, m.ngeom)), 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_geom_pair: wp.array(dtype=wp.vec2i),
nxn_pairid: wp.array(dtype=int), nxn_pairid: wp.array(dtype=int),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
# Data out: # Data out:
@@ -628,7 +628,7 @@ def _nxn_broadphase(broadphase_filter):
_add_geom_pair( _add_geom_pair(
geom_type, geom_type,
nxn_pairid, nxn_pairid,
nconmax_in, naconmax_in,
geom1, geom1,
geom2, geom2,
worldid, worldid,
@@ -668,7 +668,7 @@ def nxn_broadphase(m: Model, d: Data):
m.geom_margin, m.geom_margin,
m.nxn_geom_pair_filtered, m.nxn_geom_pair_filtered,
m.nxn_pairid_filtered, m.nxn_pairid_filtered,
d.nconmax, d.naconmax,
d.geom_xpos, d.geom_xpos,
d.geom_xmat, d.geom_xmat,
], ],
@@ -702,18 +702,18 @@ def collision(m: Model, d: Data):
distance, position, and frame. distance, position, and frame.
The results are used to populate the `d.contact` array, and the total number of contacts 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` occurred and the remaining contacts will be skipped. If this happens, raise the `nconmax`
parameter in `io.make_data` or `io.put_data`. parameter in `io.make_data` or `io.put_data`.
This function will do nothing except zero out arrays if collision detection is disabled 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 # 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 return
if m.opt.broadphase == BroadphaseType.NXN: 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 # exhaustive search over all vertices
if convex.graphadr == -1 or convex.vertnum < 10: 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) # find first support point (a)
a_dist = wp.float32(-_HUGE_VAL) max_support = wp.float32(-_HUGE_VAL)
a = wp.vec3() a = wp.vec3()
for i in range(convex.vertnum): for i in range(convex.vertnum):
vert = convex.vert[convex.vertadr + i] vert = convex.vert[convex.vertadr + i]
support = wp.dot(plane_pos_local - vert, n) support = wp.dot(plane_pos_local - vert, n)
dist = wp.where(support > threshold, support, -_HUGE_VAL) if support > max_support:
if dist > a_dist: max_support = support
indices[0] = i indices[0] = i
a_dist = dist
a = vert a = vert
if max_support < 0:
return contact_dist, contact_pos, plane_normal
threshold = max_support - 1e-3
# find point (b) furthest from a # find point (b) furthest from a
b_dist = wp.float32(-_HUGE_VAL) b_dist = wp.float32(-_HUGE_VAL)
b = wp.vec3() b = wp.vec3()
@@ -377,7 +373,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
@wp.func @wp.func
def write_contact( def write_contact(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
dist_in: float, dist_in: float,
pos_in: wp.vec3, pos_in: wp.vec3,
@@ -392,7 +388,7 @@ def write_contact(
geoms_in: wp.vec2i, geoms_in: wp.vec2i,
worldid_in: int, worldid_in: int,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -406,8 +402,8 @@ def write_contact(
contact_worldid_out: wp.array(dtype=int), contact_worldid_out: wp.array(dtype=int),
): ):
if dist_in - margin_in < 0.0: if dist_in - margin_in < 0.0:
cid = wp.atomic_add(ncon_out, 0, 1) cid = wp.atomic_add(nacon_out, 0, 1)
if cid < nconmax_in: if cid < naconmax_in:
contact_dist_out[cid] = dist_in contact_dist_out[cid] = dist_in
contact_pos_out[cid] = pos_in contact_pos_out[cid] = pos_in
contact_frame_out[cid] = frame_in contact_frame_out[cid] = frame_in
@@ -515,7 +511,7 @@ def contact_params(
@wp.func @wp.func
def plane_sphere_wrapper( def plane_sphere_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
plane: Geom, plane: Geom,
sphere: Geom, sphere: Geom,
@@ -529,7 +525,7 @@ def plane_sphere_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -547,7 +543,7 @@ def plane_sphere_wrapper(
if dist - margin < 0: if dist - margin < 0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(plane.normal), make_frame(plane.normal),
@@ -560,7 +556,7 @@ def plane_sphere_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -578,7 +574,7 @@ def plane_sphere_wrapper(
@wp.func @wp.func
def sphere_sphere_wrapper( def sphere_sphere_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
sphere1: Geom, sphere1: Geom,
sphere2: Geom, sphere2: Geom,
@@ -592,7 +588,7 @@ def sphere_sphere_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -610,7 +606,7 @@ def sphere_sphere_wrapper(
if dist - margin < 0: if dist - margin < 0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(normal), make_frame(normal),
@@ -623,7 +619,7 @@ def sphere_sphere_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -641,7 +637,7 @@ def sphere_sphere_wrapper(
@wp.func @wp.func
def sphere_capsule_wrapper( def sphere_capsule_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
sphere: Geom, sphere: Geom,
cap: Geom, cap: Geom,
@@ -655,7 +651,7 @@ def sphere_capsule_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -676,7 +672,7 @@ def sphere_capsule_wrapper(
if dist - margin < 0: if dist - margin < 0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(normal), make_frame(normal),
@@ -689,7 +685,7 @@ def sphere_capsule_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -707,7 +703,7 @@ def sphere_capsule_wrapper(
@wp.func @wp.func
def capsule_capsule_wrapper( def capsule_capsule_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
cap1: Geom, cap1: Geom,
cap2: Geom, cap2: Geom,
@@ -721,7 +717,7 @@ def capsule_capsule_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -752,7 +748,7 @@ def capsule_capsule_wrapper(
if dist - margin < 0: if dist - margin < 0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(normal), make_frame(normal),
@@ -765,7 +761,7 @@ def capsule_capsule_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -783,7 +779,7 @@ def capsule_capsule_wrapper(
@wp.func @wp.func
def plane_capsule_wrapper( def plane_capsule_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
plane: Geom, plane: Geom,
cap: Geom, cap: Geom,
@@ -797,7 +793,7 @@ def plane_capsule_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -827,7 +823,7 @@ def plane_capsule_wrapper(
disti = dist[i] disti = dist[i]
if disti - margin < 0.0: if disti - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
disti, disti,
pos[i], pos[i],
frame, frame,
@@ -840,7 +836,7 @@ def plane_capsule_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -858,7 +854,7 @@ def plane_capsule_wrapper(
@wp.func @wp.func
def plane_ellipsoid_wrapper( def plane_ellipsoid_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
plane: Geom, plane: Geom,
ellipsoid: Geom, ellipsoid: Geom,
@@ -872,7 +868,7 @@ def plane_ellipsoid_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -890,7 +886,7 @@ def plane_ellipsoid_wrapper(
if dist - margin < 0: if dist - margin < 0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(normal), make_frame(normal),
@@ -903,7 +899,7 @@ def plane_ellipsoid_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -921,7 +917,7 @@ def plane_ellipsoid_wrapper(
@wp.func @wp.func
def plane_box_wrapper( def plane_box_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
plane: Geom, plane: Geom,
box: Geom, box: Geom,
@@ -935,7 +931,7 @@ def plane_box_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -956,7 +952,7 @@ def plane_box_wrapper(
disti = dist[i] disti = dist[i]
if disti - margin < 0.0: if disti - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
disti, disti,
pos[i], pos[i],
frame, frame,
@@ -969,7 +965,7 @@ def plane_box_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -990,7 +986,7 @@ _HUGE_VAL = 1e6
@wp.func @wp.func
def plane_convex_wrapper( def plane_convex_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
plane: Geom, plane: Geom,
convex: Geom, convex: Geom,
@@ -1004,7 +1000,7 @@ def plane_convex_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1025,7 +1021,7 @@ def plane_convex_wrapper(
disti = dist[i] disti = dist[i]
if disti - margin < 0.0: if disti - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
disti, disti,
pos[i], pos[i],
frame, frame,
@@ -1038,7 +1034,7 @@ def plane_convex_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1056,7 +1052,7 @@ def plane_convex_wrapper(
@wp.func @wp.func
def sphere_cylinder_wrapper( def sphere_cylinder_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
sphere: Geom, sphere: Geom,
cylinder: Geom, cylinder: Geom,
@@ -1070,7 +1066,7 @@ def sphere_cylinder_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1098,7 +1094,7 @@ def sphere_cylinder_wrapper(
if dist - margin < 0.0: if dist - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(normal), make_frame(normal),
@@ -1111,7 +1107,7 @@ def sphere_cylinder_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1129,7 +1125,7 @@ def sphere_cylinder_wrapper(
@wp.func @wp.func
def plane_cylinder_wrapper( def plane_cylinder_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
plane: Geom, plane: Geom,
cylinder: Geom, cylinder: Geom,
@@ -1143,7 +1139,7 @@ def plane_cylinder_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1174,7 +1170,7 @@ def plane_cylinder_wrapper(
disti = dist[i] disti = dist[i]
if disti - margin < 0.0: if disti - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
disti, disti,
pos[i], pos[i],
frame, frame,
@@ -1187,7 +1183,7 @@ def plane_cylinder_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1205,7 +1201,7 @@ def plane_cylinder_wrapper(
@wp.func @wp.func
def sphere_box_wrapper( def sphere_box_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
sphere: Geom, sphere: Geom,
box: Geom, box: Geom,
@@ -1219,7 +1215,7 @@ def sphere_box_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1236,7 +1232,7 @@ def sphere_box_wrapper(
if dist - margin < 0.0: if dist - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(normal), make_frame(normal),
@@ -1249,7 +1245,7 @@ def sphere_box_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1267,7 +1263,7 @@ def sphere_box_wrapper(
@wp.func @wp.func
def capsule_box_wrapper( def capsule_box_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
cap: Geom, cap: Geom,
box: Geom, box: Geom,
@@ -1281,7 +1277,7 @@ def capsule_box_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1314,7 +1310,7 @@ def capsule_box_wrapper(
disti = dist[i] disti = dist[i]
if disti - margin < 0.0: if disti - margin < 0.0:
write_contact( write_contact(
nconmax_in, naconmax_in,
disti, disti,
pos[i], pos[i],
make_frame(normal[i]), make_frame(normal[i]),
@@ -1327,7 +1323,7 @@ def capsule_box_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1345,7 +1341,7 @@ def capsule_box_wrapper(
@wp.func @wp.func
def box_box_wrapper( def box_box_wrapper(
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
# In: # In:
box1: Geom, box1: Geom,
box2: Geom, box2: Geom,
@@ -1359,7 +1355,7 @@ def box_box_wrapper(
solimp: vec5, solimp: vec5,
geoms: wp.vec2i, geoms: wp.vec2i,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1388,7 +1384,7 @@ def box_box_wrapper(
continue continue
write_contact( write_contact(
nconmax_in, naconmax_in,
dist[i], dist[i],
pos[i], pos[i],
make_frame(normal[i]), make_frame(normal[i]),
@@ -1401,7 +1397,7 @@ def box_box_wrapper(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1446,12 +1442,17 @@ def _check_primitive_collisions():
assert _check_primitive_collisions(), "_PRIMITIVE_COLLISIONS is in invalid order" 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(enable_backward=False)
@nested_kernel(module="unique", enable_backward=False)
def _primitive_narrowphase( def _primitive_narrowphase(
# Model: # Model:
geom_type: wp.array(dtype=int), geom_type: wp.array(dtype=int),
@@ -1492,7 +1493,7 @@ def _create_narrowphase_kernel():
pair_gap: wp.array2d(dtype=float), pair_gap: wp.array2d(dtype=float),
pair_friction: wp.array2d(dtype=vec5), pair_friction: wp.array2d(dtype=vec5),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array(dtype=wp.vec2i),
@@ -1500,7 +1501,7 @@ def _create_narrowphase_kernel():
collision_worldid_in: wp.array(dtype=int), collision_worldid_in: wp.array(dtype=int),
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array(dtype=int),
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -1595,13 +1596,13 @@ def _create_narrowphase_kernel():
geom_xmat_in[worldid, g2], geom_xmat_in[worldid, g2],
) )
for i in range(wp.static(len(_primitive_collisions_func))): for i in range(wp.static(len(primitive_collisions_func))):
collision_type1 = wp.static(_primitive_collisions_types[i][0]) collision_type1 = wp.static(primitive_collisions_types[i][0])
collision_type2 = wp.static(_primitive_collisions_types[i][1]) collision_type2 = wp.static(primitive_collisions_types[i][1])
if collision_type1 == type1 and collision_type2 == type2: if collision_type1 == type1 and collision_type2 == type2:
wp.static(_primitive_collisions_func[i])( wp.static(primitive_collisions_func[i])(
nconmax_in, naconmax_in,
geom1, geom1,
geom2, geom2,
worldid, worldid,
@@ -1613,7 +1614,7 @@ def _create_narrowphase_kernel():
solreffriction, solreffriction,
solimp, solimp,
geoms, geoms,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -1631,13 +1632,16 @@ def _create_narrowphase_kernel():
def _primitive_narrowphase_builder(m: Model): def _primitive_narrowphase_builder(m: Model):
_primitive_collisions_types = []
_primitive_collisions_func = []
for types, func in _PRIMITIVE_COLLISIONS.items(): for types, func in _PRIMITIVE_COLLISIONS.items():
idx = upper_trid_index(len(GeomType), types[0].value, types[1].value) 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: if m.geom_pair_type_count[idx] and types not in _primitive_collisions_types:
_primitive_collisions_types.append(types) _primitive_collisions_types.append(types)
_primitive_collisions_func.append(func) _primitive_collisions_func.append(func)
return _create_narrowphase_kernel() return _create_narrowphase_kernel(_primitive_collisions_types, _primitive_collisions_func)
@event_scope @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. # for pair types without collisions, as well as updating the launch dimensions.
wp.launch( wp.launch(
_primitive_narrowphase_builder(m), _primitive_narrowphase_builder(m),
dim=d.nconmax, dim=d.naconmax,
inputs=[ inputs=[
m.geom_type, m.geom_type,
m.geom_condim, m.geom_condim,
@@ -1699,7 +1703,7 @@ def primitive_narrowphase(m: Model, d: Data):
m.pair_margin, m.pair_margin,
m.pair_gap, m.pair_gap,
m.pair_friction, m.pair_friction,
d.nconmax, d.naconmax,
d.geom_xpos, d.geom_xpos,
d.geom_xmat, d.geom_xmat,
d.collision_pair, d.collision_pair,
@@ -1708,7 +1712,7 @@ def primitive_narrowphase(m: Model, d: Data):
d.ncollision, d.ncollision,
], ],
outputs=[ outputs=[
d.ncon, d.nacon,
d.contact.dist, d.contact.dist,
d.contact.pos, d.contact.pos,
d.contact.frame, 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), plugin_attr: wp.array(dtype=wp.vec3f),
geom_plugin_index: wp.array(dtype=int), geom_plugin_index: wp.array(dtype=int),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
collision_pair_in: wp.array(dtype=wp.vec2i), collision_pair_in: wp.array(dtype=wp.vec2i),
@@ -679,7 +679,7 @@ def _sdf_narrowphase(
sdf_initpoints: int, sdf_initpoints: int,
sdf_iterations: int, sdf_iterations: int,
# Data out: # Data out:
ncon_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int),
contact_dist_out: wp.array(dtype=float), contact_dist_out: wp.array(dtype=float),
contact_pos_out: wp.array(dtype=wp.vec3), contact_pos_out: wp.array(dtype=wp.vec3),
contact_frame_out: wp.array(dtype=wp.mat33), contact_frame_out: wp.array(dtype=wp.mat33),
@@ -850,7 +850,7 @@ def _sdf_narrowphase(
mesh_data2, mesh_data2,
) )
write_contact( write_contact(
nconmax_in, naconmax_in,
dist, dist,
pos, pos,
make_frame(n), make_frame(n),
@@ -863,7 +863,7 @@ def _sdf_narrowphase(
solimp, solimp,
geoms, geoms,
worldid, worldid,
ncon_out, nacon_out,
contact_dist_out, contact_dist_out,
contact_pos_out, contact_pos_out,
contact_frame_out, contact_frame_out,
@@ -882,7 +882,7 @@ def _sdf_narrowphase(
def sdf_narrowphase(m: Model, d: Data): def sdf_narrowphase(m: Model, d: Data):
wp.launch( wp.launch(
_sdf_narrowphase, _sdf_narrowphase,
dim=(m.opt.sdf_initpoints, d.nconmax), dim=(m.opt.sdf_initpoints, d.naconmax),
inputs=[ inputs=[
m.nmeshface, m.nmeshface,
m.geom_type, m.geom_type,
@@ -931,7 +931,7 @@ def sdf_narrowphase(m: Model, d: Data):
m.plugin, m.plugin,
m.plugin_attr, m.plugin_attr,
m.geom_plugin_index, m.geom_plugin_index,
d.nconmax, d.naconmax,
d.geom_xpos, d.geom_xpos,
d.geom_xmat, d.geom_xmat,
d.collision_pair, d.collision_pair,
@@ -942,7 +942,7 @@ def sdf_narrowphase(m: Model, d: Data):
m.opt.sdf_iterations, m.opt.sdf_iterations,
], ],
outputs=[ outputs=[
d.ncon, d.nacon,
d.contact.dist, d.contact.dist,
d.contact.pos, d.contact.pos,
d.contact.frame, 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), geom_bodyid: wp.array(dtype=int),
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d(dtype=float),
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d(dtype=wp.spatial_vector),
@@ -1141,7 +1141,7 @@ def _efc_contact_pyramidal(
): ):
conid, dimid = wp.tid() conid, dimid = wp.tid()
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
condim = condim_in[conid] condim = condim_in[conid]
@@ -1275,7 +1275,7 @@ def _efc_contact_elliptic(
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int),
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
qvel_in: wp.array2d(dtype=float), qvel_in: wp.array2d(dtype=float),
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d(dtype=wp.spatial_vector),
@@ -1307,7 +1307,7 @@ def _efc_contact_elliptic(
): ):
conid, dimid = wp.tid() conid, dimid = wp.tid()
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
condim = condim_in[conid] condim = condim_in[conid]
@@ -1799,7 +1799,7 @@ def make_constraint(m: types.Model, d: types.Data):
if m.opt.cone == types.ConeType.PYRAMIDAL: if m.opt.cone == types.ConeType.PYRAMIDAL:
wp.launch( wp.launch(
_efc_contact_pyramidal, _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=[ inputs=[
m.nv, m.nv,
m.opt.timestep, m.opt.timestep,
@@ -1810,7 +1810,7 @@ def make_constraint(m: types.Model, d: types.Data):
m.dof_bodyid, m.dof_bodyid,
m.geom_bodyid, m.geom_bodyid,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.qvel, d.qvel,
d.subtree_com, d.subtree_com,
d.cdof, d.cdof,
@@ -1843,7 +1843,7 @@ def make_constraint(m: types.Model, d: types.Data):
elif m.opt.cone == types.ConeType.ELLIPTIC: elif m.opt.cone == types.ConeType.ELLIPTIC:
wp.launch( wp.launch(
_efc_contact_elliptic, _efc_contact_elliptic,
dim=(d.nconmax, m.condim_max), dim=(d.naconmax, m.condim_max),
inputs=[ inputs=[
m.nv, m.nv,
m.opt.timestep, m.opt.timestep,
@@ -1854,7 +1854,7 @@ def make_constraint(m: types.Model, d: types.Data):
m.dof_bodyid, m.dof_bodyid,
m.geom_bodyid, m.geom_bodyid,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.qvel, d.qvel,
d.subtree_com, d.subtree_com,
d.cdof, 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: # Model:
opt_timestep: wp.array(dtype=float), opt_timestep: wp.array(dtype=float),
# Data in: # Data in:
nconmax_in: int, nworld_in: int,
naconmax_in: int,
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int), nefc_in: wp.array(dtype=int),
time_in: wp.array(dtype=float), time_in: wp.array(dtype=float),
ncollision_in: wp.array(dtype=int), ncollision_in: wp.array(dtype=int),
@@ -224,11 +225,13 @@ def _next_time(
if worldid == 0: if worldid == 0:
ncollision = ncollision_in[0] ncollision = ncollision_in[0]
if ncollision > nconmax_in: if ncollision > naconmax_in:
wp.printf("ncollision overflow - please increase nconmax to %u\n", ncollision) 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: if nacon_in[0] > naconmax_in:
wp.printf("ncon overflow - please increase nconmax to %u\n", ncon_in[0]) 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): 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,), dim=(d.nworld,),
inputs=[ inputs=[
m.opt.timestep, m.opt.timestep,
d.nconmax, d.nworld,
d.naconmax,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.nefc, d.nefc,
d.time, d.time,
d.ncollision, 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(): if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any():
raise NotImplementedError("Flex collisions are not implemented.") raise NotImplementedError("Flex collisions are not implemented.")
if mjm.geom_fluid.any():
raise NotImplementedError("Ellipsoid fluid model not implemented.")
# check options # check options
for opt, opt_types, msg in ( for opt, opt_types, msg in (
(mjm.opt.integrator, types.IntegratorType, "Integrator"), (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.", "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( m = types.Model(
nq=mjm.nq, nq=mjm.nq,
nv=mjm.nv, 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_contype=wp.array(mjm.body_contype, dtype=int),
body_conaffinity=wp.array(mjm.body_conaffinity, dtype=int), body_conaffinity=wp.array(mjm.body_conaffinity, dtype=int),
body_gravcomp=create_nmodel_batched_array(mjm.body_gravcomp, dtype=float), 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_type=wp.array(mjm.jnt_type, dtype=int),
jnt_qposadr=wp.array(mjm.jnt_qposadr, dtype=int), jnt_qposadr=wp.array(mjm.jnt_qposadr, dtype=int),
jnt_dofadr=wp.array(mjm.jnt_dofadr, 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_solref=create_nmodel_batched_array(mjm.geom_solref, dtype=wp.vec2),
geom_solimp=create_nmodel_batched_array(mjm.geom_solimp, dtype=types.vec5), 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_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_aabb=wp.array2d(mjm.geom_aabb, dtype=wp.vec3),
geom_rbound=create_nmodel_batched_array(mjm.geom_rbound, dtype=float), geom_rbound=create_nmodel_batched_array(mjm.geom_rbound, dtype=float),
geom_pos=create_nmodel_batched_array(mjm.geom_pos, dtype=wp.vec3), 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 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. Creates a data object on device.
Args: Args:
mjm (mujoco.MjModel): The model containing kinematic and dynamic information (host). mjm (mujoco.MjModel): The model containing kinematic and dynamic information (host).
nworld (int, optional): Number of worlds. Defaults to 1. nworld (int, optional): Number of worlds. Defaults to 1.
nconmax (int, optional): Maximum number of contacts for all worlds. Defaults to -1. nworld (int, optional): The number of worlds. Defaults to 1.
njmax (int, optional): 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: Returns:
Data: The data object containing the current state and output arrays (device). Data: The data object containing the current state and output arrays (device).
""" """
# TODO(team): move to Model? # TODO(team): move nconmax, njmax to Model?
if nconmax == -1: # TODO(team): improve heuristic for nconmax and njmax
# TODO(team): heuristic for nconmax nconmax = nconmax or 20
nconmax = nworld * 20 njmax = njmax or nconmax * 6
if njmax == -1:
# TODO(team): heuristic for njmax
njmax = 20 * 6
if nworld < 1 or nworld > MAX_WORLDS: if nworld < 1 or nworld > MAX_WORLDS:
raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}") raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}")
if nconmax < 0: if naconmax is None:
raise ValueError("nconmax must be >= 0") 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: if njmax < 0:
raise ValueError("njmax must be >= 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): if mujoco.mj_isSparse(mjm):
qM = wp.zeros((nworld, 1, mjm.nM), dtype=float) qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
qLD = wp.zeros((nworld, 1, mjm.nC), 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) qM_integration = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
qLD_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) nsensorcontact = np.sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT)
nrangefinder = sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER) nrangefinder = sum(mjm.sensor_type == mujoco.mjtSensor.mjSENS_RANGEFINDER)
return types.Data( return types.Data(
nworld=nworld, nworld=nworld,
nconmax=nconmax, naconmax=naconmax,
njmax=njmax, njmax=njmax,
solver_niter=wp.zeros(nworld, dtype=int), 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=wp.zeros(nworld, dtype=int),
ne_connect=wp.zeros(nworld, dtype=int), # warp only ne_connect=wp.zeros(nworld, dtype=int), # warp only
ne_weld=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_constraint=wp.zeros((nworld, mjm.nv), dtype=float),
qfrc_inverse=wp.zeros((nworld, mjm.nv), dtype=float), qfrc_inverse=wp.zeros((nworld, mjm.nv), dtype=float),
contact=types.Contact( contact=types.Contact(
dist=wp.zeros((nconmax,), dtype=float), dist=wp.zeros((naconmax,), dtype=float),
pos=wp.zeros((nconmax,), dtype=wp.vec3f), pos=wp.zeros((naconmax,), dtype=wp.vec3f),
frame=wp.zeros((nconmax,), dtype=wp.mat33f), frame=wp.zeros((naconmax,), dtype=wp.mat33f),
includemargin=wp.zeros((nconmax,), dtype=float), includemargin=wp.zeros((naconmax,), dtype=float),
friction=wp.zeros((nconmax,), dtype=types.vec5), friction=wp.zeros((naconmax,), dtype=types.vec5),
solref=wp.zeros((nconmax,), dtype=wp.vec2f), solref=wp.zeros((naconmax,), dtype=wp.vec2f),
solreffriction=wp.zeros((nconmax,), dtype=wp.vec2f), solreffriction=wp.zeros((naconmax,), dtype=wp.vec2f),
solimp=wp.zeros((nconmax,), dtype=types.vec5), solimp=wp.zeros((naconmax,), dtype=types.vec5),
dim=wp.zeros((nconmax,), dtype=int), dim=wp.zeros((naconmax,), dtype=int),
geom=wp.zeros((nconmax,), dtype=wp.vec2i), geom=wp.zeros((naconmax,), dtype=wp.vec2i),
efc_address=wp.zeros( efc_address=wp.zeros((naconmax, np.maximum(1, 2 * (condim_max - 1))), dtype=int),
(nconmax, np.maximum(1, 2 * (condim_max - 1))), worldid=wp.zeros((naconmax,), dtype=int),
dtype=int,
),
worldid=wp.zeros((nconmax,), dtype=int),
), ),
efc=types.Constraint( efc=types.Constraint(
type=wp.zeros((nworld, njmax), dtype=int), 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 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 driver
collision_pair=wp.zeros((nconmax,), dtype=wp.vec2i), collision_pair=wp.zeros((naconmax,), dtype=wp.vec2i),
collision_pairid=wp.zeros((nconmax,), dtype=int), collision_pairid=wp.zeros((naconmax,), dtype=int),
collision_worldid=wp.zeros((nconmax,), dtype=int), collision_worldid=wp.zeros((naconmax,), dtype=int),
ncollision=wp.zeros((1,), dtype=int), ncollision=wp.zeros((1,), dtype=int),
# narrowphase (EPA polytope) # narrowphase (EPA polytope)
epa_vert=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3), epa_vert=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert1=wp.zeros(shape=(nconmax, 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=(nconmax, 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=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=int), epa_vert_index1=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_vert_index2=wp.zeros(shape=(nconmax, 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=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3i), epa_face=wp.zeros(shape=(naconmax, 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_pr=wp.zeros(shape=(naconmax, 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_norm2=wp.zeros(shape=(naconmax, 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_index=wp.zeros(shape=(naconmax, 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_map=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_horizon=wp.zeros(shape=(nconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int), epa_horizon=wp.zeros(shape=(naconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int),
multiccd_polygon=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3), multiccd_polygon=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_clipped=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3), multiccd_clipped=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_pnormal=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3), multiccd_pnormal=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_pdist=wp.zeros(shape=(nconmax, max_npolygon), dtype=float), multiccd_pdist=wp.zeros(shape=(naconmax, max_npolygon), dtype=float),
multiccd_idx1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int), multiccd_idx1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_idx2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int), multiccd_idx2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_n1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3), multiccd_n1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_n2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3), multiccd_n2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_endvert=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3), multiccd_endvert=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_face1=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3), multiccd_face1=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_face2=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3), multiccd_face2=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
# rne_postconstraint # rne_postconstraint
cacc=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), cacc=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector),
cfrc_int=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( def put_data(
mjm: mujoco.MjModel, mjm: mujoco.MjModel,
mjd: mujoco.MjData, mjd: mujoco.MjData,
nworld: Optional[int] = None, nworld: int = 1,
nconmax: Optional[int] = None, nconmax: Optional[int] = None,
njmax: Optional[int] = None, njmax: Optional[int] = None,
naconmax: Optional[int] = None,
) -> types.Data: ) -> types.Data:
""" """
Moves data from host to a device. 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). mjm (mujoco.MjModel): The model containing kinematic and dynamic information (host).
mjd (mujoco.MjData): The data object containing current state and output arrays (host). mjd (mujoco.MjData): The data object containing current state and output arrays (host).
nworld (int, optional): The number of worlds. Defaults to 1. nworld (int, optional): The number of worlds. Defaults to 1.
nconmax (int, optional): The maximum number of contacts for all worlds. Defaults to -1. nconmax (int, optional): Number of contacts to allocate per world. Contacts exist in large
njmax (int, optional): The maximum number of constraints per world. Defaults to -1. 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: Returns:
Data: The data object containing the current state and output arrays (device). 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 # 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 # we need to ensure these are only workspace fields and don't carry state
nworld = nworld or 1 # TODO(team): better heuristic for nconmax and njmax
# TODO(team): better heuristic for nconmax nconmax = nconmax or max(5, 4 * mjd.ncon)
nconmax = nconmax or max(512, 4 * mjd.ncon * nworld)
# TODO(team): better heuristic for njmax
njmax = njmax or max(5, 4 * mjd.nefc) njmax = njmax or max(5, 4 * mjd.nefc)
if nworld < 1 or nworld > MAX_WORLDS: if nworld < 1 or nworld > MAX_WORLDS:
raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}") raise ValueError(f"nworld must be >= 1 and <= {MAX_WORLDS}")
if nconmax < 0: if naconmax is None:
raise ValueError("nconmax must be >= 0") 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: if njmax < 0:
raise ValueError("njmax must be >= 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: if mjd.nefc > njmax:
raise ValueError(f"njmax overflow (njmax must be >= {mjd.nefc})") 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 = np.concatenate((mjm.geom_condim, mjm.pair_dim))
condim_max = np.max(condim) if len(condim) > 0 else 0 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 i in range(nworld):
for j in range(mjd.ncon): for j in range(mjd.ncon):
condim = mjd.contact.dim[j] condim = mjd.contact.dim[j]
@@ -1275,7 +1298,7 @@ def put_data(
for k in range(nconvar): for k in range(nconvar):
contact_efc_address[i * mjd.ncon + j, k] = mjd.nefc * i + efc_address + k 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_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)) 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( return types.Data(
nworld=nworld, nworld=nworld,
nconmax=nconmax, naconmax=naconmax,
njmax=njmax, njmax=njmax,
solver_niter=tile(mjd.solver_niter[0]), 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=wp.full(shape=(nworld), value=mjd.ne),
ne_connect=wp.full(shape=(nworld), value=ne_connect), ne_connect=wp.full(shape=(nworld), value=ne_connect),
ne_weld=wp.full(shape=(nworld), value=ne_weld), ne_weld=wp.full(shape=(nworld), value=ne_weld),
@@ -1412,16 +1435,16 @@ def put_data(
qfrc_constraint=tile(mjd.qfrc_constraint), qfrc_constraint=tile(mjd.qfrc_constraint),
qfrc_inverse=tile(mjd.qfrc_inverse), qfrc_inverse=tile(mjd.qfrc_inverse),
contact=types.Contact( contact=types.Contact(
dist=padtile(mjd.contact.dist, nconmax), dist=padtile(mjd.contact.dist, naconmax),
pos=padtile(mjd.contact.pos, nconmax, dtype=wp.vec3), pos=padtile(mjd.contact.pos, naconmax, dtype=wp.vec3),
frame=padtile(mjd.contact.frame, nconmax, dtype=wp.mat33), frame=padtile(mjd.contact.frame, naconmax, dtype=wp.mat33),
includemargin=padtile(mjd.contact.includemargin, nconmax), includemargin=padtile(mjd.contact.includemargin, naconmax),
friction=padtile(mjd.contact.friction, nconmax, dtype=types.vec5), friction=padtile(mjd.contact.friction, naconmax, dtype=types.vec5),
solref=padtile(mjd.contact.solref, nconmax, dtype=wp.vec2f), solref=padtile(mjd.contact.solref, naconmax, dtype=wp.vec2f),
solreffriction=padtile(mjd.contact.solreffriction, nconmax, dtype=wp.vec2f), solreffriction=padtile(mjd.contact.solreffriction, naconmax, dtype=wp.vec2f),
solimp=padtile(mjd.contact.solimp, nconmax, dtype=types.vec5), solimp=padtile(mjd.contact.solimp, naconmax, dtype=types.vec5),
dim=padtile(mjd.contact.dim, nconmax), dim=padtile(mjd.contact.dim, naconmax),
geom=padtile(mjd.contact.geom, nconmax, dtype=wp.vec2i), geom=padtile(mjd.contact.geom, naconmax, dtype=wp.vec2i),
efc_address=arr(contact_efc_address), efc_address=arr(contact_efc_address),
worldid=arr(contact_worldid), worldid=arr(contact_worldid),
), ),
@@ -1483,33 +1506,33 @@ def put_data(
sap_cumulative_sum=wp.zeros((nworld, mjm.ngeom), dtype=int), 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))), 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 driver
collision_pair=wp.empty(nconmax, dtype=wp.vec2i), collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
collision_pairid=wp.empty(nconmax, dtype=int), collision_pairid=wp.empty(naconmax, dtype=int),
collision_worldid=wp.empty(nconmax, dtype=int), collision_worldid=wp.empty(naconmax, dtype=int),
ncollision=wp.zeros(1, dtype=int), ncollision=wp.zeros(1, dtype=int),
# narrowphase (EPA polytope) # narrowphase (EPA polytope)
epa_vert=wp.zeros(shape=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3), epa_vert=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=wp.vec3),
epa_vert1=wp.zeros(shape=(nconmax, 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=(nconmax, 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=(nconmax, 5 + mjm.opt.ccd_iterations), dtype=int), epa_vert_index1=wp.zeros(shape=(naconmax, 5 + mjm.opt.ccd_iterations), dtype=int),
epa_vert_index2=wp.zeros(shape=(nconmax, 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=(nconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=wp.vec3i), epa_face=wp.zeros(shape=(naconmax, 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_pr=wp.zeros(shape=(naconmax, 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_norm2=wp.zeros(shape=(naconmax, 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_index=wp.zeros(shape=(naconmax, 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_map=wp.zeros(shape=(naconmax, 6 + types.MJ_MAX_EPAFACES * mjm.opt.ccd_iterations), dtype=int),
epa_horizon=wp.zeros(shape=(nconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int), epa_horizon=wp.zeros(shape=(naconmax, 2 * types.MJ_MAX_EPAHORIZON), dtype=int),
multiccd_polygon=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3), multiccd_polygon=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_clipped=wp.zeros(shape=(nconmax, 2 * max_npolygon), dtype=wp.vec3), multiccd_clipped=wp.zeros(shape=(naconmax, 2 * max_npolygon), dtype=wp.vec3),
multiccd_pnormal=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3), multiccd_pnormal=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_pdist=wp.zeros(shape=(nconmax, max_npolygon), dtype=float), multiccd_pdist=wp.zeros(shape=(naconmax, max_npolygon), dtype=float),
multiccd_idx1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int), multiccd_idx1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_idx2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=int), multiccd_idx2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=int),
multiccd_n1=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3), multiccd_n1=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_n2=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3), multiccd_n2=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_endvert=wp.zeros(shape=(nconmax, max_meshdegree), dtype=wp.vec3), multiccd_endvert=wp.zeros(shape=(naconmax, max_meshdegree), dtype=wp.vec3),
multiccd_face1=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3), multiccd_face1=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
multiccd_face2=wp.zeros(shape=(nconmax, max_npolygon), dtype=wp.vec3), multiccd_face2=wp.zeros(shape=(naconmax, max_npolygon), dtype=wp.vec3),
# rne_postconstraint but also smooth # rne_postconstraint but also smooth
cacc=tile(mjd.cacc, dtype=wp.spatial_vector), cacc=tile(mjd.cacc, dtype=wp.spatial_vector),
cfrc_int=tile(mjd.cfrc_int, 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] result.solver_niter[0] = d.solver_niter.numpy()[0]
ncon = d.ncon.numpy()[0] nacon = d.nacon.numpy()[0]
nefc = d.nefc.numpy()[0] nefc = d.nefc.numpy()[0]
if ncon != result.ncon or nefc != result.nefc: if nacon != result.ncon or nefc != result.nefc:
mujoco._functions._realloc_con_efc(result, ncon=ncon, nefc=nefc) mujoco._functions._realloc_con_efc(result, ncon=nacon, nefc=nefc)
result.time = d.time.numpy()[0] result.time = d.time.numpy()[0]
result.energy = d.energy.numpy()[0] result.energy = d.energy.numpy()[0]
@@ -1635,16 +1658,16 @@ def get_data_into(
result.act = d.act.numpy()[0] result.act = d.act.numpy()[0]
result.act_dot = d.act_dot.numpy()[0] result.act_dot = d.act_dot.numpy()[0]
result.contact.dist[:] = d.contact.dist.numpy()[:ncon] result.contact.dist[:] = d.contact.dist.numpy()[:nacon]
result.contact.pos[:] = d.contact.pos.numpy()[:ncon] result.contact.pos[:] = d.contact.pos.numpy()[:nacon]
result.contact.frame[:] = d.contact.frame.numpy()[:ncon].reshape((-1, 9)) result.contact.frame[:] = d.contact.frame.numpy()[:nacon].reshape((-1, 9))
result.contact.includemargin[:] = d.contact.includemargin.numpy()[:ncon] result.contact.includemargin[:] = d.contact.includemargin.numpy()[:nacon]
result.contact.friction[:] = d.contact.friction.numpy()[:ncon] result.contact.friction[:] = d.contact.friction.numpy()[:nacon]
result.contact.solref[:] = d.contact.solref.numpy()[:ncon] result.contact.solref[:] = d.contact.solref.numpy()[:nacon]
result.contact.solreffriction[:] = d.contact.solreffriction.numpy()[:ncon] result.contact.solreffriction[:] = d.contact.solreffriction.numpy()[:nacon]
result.contact.solimp[:] = d.contact.solimp.numpy()[:ncon] result.contact.solimp[:] = d.contact.solimp.numpy()[:nacon]
result.contact.dim[:] = d.contact.dim.numpy()[:ncon] result.contact.dim[:] = d.contact.dim.numpy()[:nacon]
result.contact.efc_address[:] = d.contact.efc_address.numpy()[:ncon, 0] result.contact.efc_address[:] = d.contact.efc_address.numpy()[:nacon, 0]
if mujoco.mj_isSparse(mjm): if mujoco.mj_isSparse(mjm):
result.qM[:] = d.qM.numpy()[0, 0] result.qM[:] = d.qM.numpy()[0, 0]
@@ -1704,8 +1727,43 @@ def get_data_into(
result.sensordata[:] = d.sensordata.numpy() result.sensordata[:] = d.sensordata.numpy()
# TODO(thowell): shared @wp.func for _reset kernel?
@wp.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: # Model:
nq: int, nq: int,
nv: int, nv: int,
@@ -1719,7 +1777,7 @@ def _reset_nworld(
nworld_in: int, nworld_in: int,
# Data out: # Data out:
solver_niter_out: wp.array(dtype=int), 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_out: wp.array(dtype=int),
ne_connect_out: wp.array(dtype=int), ne_connect_out: wp.array(dtype=int),
ne_weld_out: wp.array(dtype=int), ne_weld_out: wp.array(dtype=int),
@@ -1746,7 +1804,7 @@ def _reset_nworld(
solver_niter_out[worldid] = 0 solver_niter_out[worldid] = 0
if worldid == 0: if worldid == 0:
ncon_out[0] = 0 nacon_out[0] = 0
ne_out[worldid] = 0 ne_out[worldid] = 0
ne_connect_out[worldid] = 0 ne_connect_out[worldid] = 0
ne_weld_out[worldid] = 0 ne_weld_out[worldid] = 0
@@ -1778,7 +1836,85 @@ def _reset_nworld(
@wp.kernel @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: # Model:
body_mocapid: wp.array(dtype=int), body_mocapid: wp.array(dtype=int),
body_pos: wp.array2d(dtype=wp.vec3), body_pos: wp.array2d(dtype=wp.vec3),
@@ -1797,9 +1933,33 @@ def _reset_mocap(
@wp.kernel @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: # Data in:
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
# In: # In:
nefcaddress: int, nefcaddress: int,
# Data out: # Data out:
@@ -1818,9 +1978,55 @@ def _reset_contact(
): ):
conid = wp.tid() conid = wp.tid()
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return 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_dist_out[conid] = 0.0
contact_pos_out[conid] = wp.vec3(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) 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 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.""" """Clear data, set defaults."""
d.xfrc_applied.zero_() if m.opt.is_sparse:
d.qM.zero_() qM_dim = (1, m.nM)
else:
qM_dim = (m.nv, m.nv)
# set mocap_pos/quat = body_pos/quat for mocap bodies if reset is not None:
wp.launch( wp.launch(_reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset], outputs=[d.xfrc_applied])
_reset_mocap, 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_qM, dim=(d.nworld, qM_dim[0], qM_dim[1]), inputs=[reset], outputs=[d.qM])
)
# clear contacts # set mocap_pos/quat = body_pos/quat for mocap bodies
wp.launch( wp.launch(
_reset_contact, _reset_mocap,
dim=d.nconmax, dim=(d.nworld, m.nbody),
inputs=[d.ncon, d.contact.efc_address.shape[1]], inputs=[m.body_mocapid, m.body_pos, m.body_quat, reset],
outputs=[ outputs=[d.mocap_pos, d.mocap_quat],
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( # clear contacts
_reset_nworld, wp.launch(
dim=d.nworld, _reset_contact,
inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld], dim=d.naconmax,
outputs=[ inputs=[d.nacon, reset, d.contact.efc_address.shape[1]],
d.solver_niter, outputs=[
d.ncon, d.contact.dist,
d.ne, d.contact.pos,
d.ne_connect, d.contact.frame,
d.ne_weld, d.contact.includemargin,
d.ne_jnt, d.contact.friction,
d.ne_ten, d.contact.solref,
d.nf, d.contact.solreffriction,
d.nl, d.contact.solimp,
d.nefc, d.contact.dim,
d.nsolving, d.contact.geom,
d.time, d.contact.efc_address,
d.energy, d.contact.worldid,
d.qpos, ],
d.qvel, )
d.act,
d.qacc_warmstart, wp.launch(
d.ctrl, _reset_nworld,
d.qfrc_applied, dim=d.nworld,
d.eq_active, inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld, reset],
d.qacc, outputs=[
d.act_dot, d.solver_niter,
d.sensordata, 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]]): 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 MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data 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 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 JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope 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.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 @wp.kernel
def _spring_damper_dof_passive( def _spring_damper_dof_passive(
# Model: # Model:
@@ -224,50 +264,185 @@ def _gravity_force(
@wp.kernel @wp.kernel
def _box_fluid( def _fluid_force(
# Model: # Model:
opt_wind: wp.array(dtype=wp.vec3), opt_wind: wp.array(dtype=wp.vec3),
opt_density: wp.array(dtype=float), opt_density: wp.array(dtype=float),
opt_viscosity: wp.array(dtype=float), opt_viscosity: wp.array(dtype=float),
body_rootid: wp.array(dtype=int), 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_mass: wp.array2d(dtype=float),
body_inertia: wp.array2d(dtype=wp.vec3), 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: # Data in:
xipos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33), 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), subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector), cvel_in: wp.array2d(dtype=wp.spatial_vector),
# Data out: # Data out:
fluid_applied_out: wp.array2d(dtype=wp.spatial_vector), 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() 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] wind = opt_wind[worldid]
density = opt_density[worldid] density = opt_density[worldid]
viscosity = opt_viscosity[worldid] viscosity = opt_viscosity[worldid]
# map from CoM-centered to local body-centered 6D velocity # Body kinematics
xipos = xipos_in[worldid, bodyid]
# body-inertial
pos = xipos_in[worldid, bodyid]
rot = ximat_in[worldid, bodyid] rot = ximat_in[worldid, bodyid]
rotT = wp.transpose(rot) rotT = wp.transpose(rot)
# transform velocity
cvel = cvel_in[worldid, bodyid] cvel = cvel_in[worldid, bodyid]
torque = wp.spatial_top(cvel) ang_global = wp.spatial_top(cvel)
force = wp.spatial_bottom(cvel) lin_global = wp.spatial_bottom(cvel)
subtree_com = subtree_com_in[worldid, body_rootid[bodyid]] subtree_root = subtree_com_in[worldid, body_rootid[bodyid]]
dif = pos - subtree_com lin_com = lin_global - wp.cross(xipos - subtree_root, ang_global)
force -= wp.cross(dif, torque)
lvel_torque = rotT @ torque if body_fluid_ellipsoid[bodyid]:
lvel_force = rotT @ force 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]: if wind[0] or wind[1] or wind[2]:
# subtract translational component from body velocity l_lin -= rotT @ wind
lvel_force -= rotT @ wind
lfrc_torque = wp.vec3(0.0) lfrc_torque = wp.vec3(0.0)
lfrc_force = 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) box2 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[1] - inertia[2]) * scl)
if has_viscosity: if has_viscosity:
# diameter of sphere approximation
diam = (box0 + box1 + box2) / 3.0 diam = (box0 + box1 + box2) / 3.0
lfrc_torque = -l_ang * wp.pow(diam, 3.0) * wp.pi * viscosity
# angular viscosity lfrc_force = -3.0 * l_lin * diam * wp.pi * 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
if has_density: if has_density:
# force
lfrc_force -= wp.vec3( lfrc_force -= wp.vec3(
0.5 * density * box1 * box2 * wp.abs(lvel_force[0]) * lvel_force[0], 0.5 * density * box1 * box2 * wp.abs(l_lin[0]) * l_lin[0],
0.5 * density * box0 * box2 * wp.abs(lvel_force[1]) * lvel_force[1], 0.5 * density * box0 * box2 * wp.abs(l_lin[1]) * l_lin[1],
0.5 * density * box0 * box1 * wp.abs(lvel_force[2]) * lvel_force[2], 0.5 * density * box0 * box1 * wp.abs(l_lin[2]) * l_lin[2],
) )
# torque
scl = density / 64.0 scl = density / 64.0
box0_pow4 = wp.pow(box0, 4.0) box0_pow4 = wp.pow(box0, 4.0)
box1_pow4 = wp.pow(box1, 4.0) box1_pow4 = wp.pow(box1, 4.0)
box2_pow4 = wp.pow(box2, 4.0) box2_pow4 = wp.pow(box2, 4.0)
lfrc_torque -= wp.vec3( lfrc_torque -= wp.vec3(
box0 * (box1_pow4 + box2_pow4) * wp.abs(lvel_torque[0]) * lvel_torque[0] * scl, box0 * (box1_pow4 + box2_pow4) * wp.abs(l_ang[0]) * l_ang[0] * scl,
box1 * (box0_pow4 + box2_pow4) * wp.abs(lvel_torque[1]) * lvel_torque[1] * scl, box1 * (box0_pow4 + box2_pow4) * wp.abs(l_ang[1]) * l_ang[1] * scl,
box2 * (box0_pow4 + box1_pow4) * wp.abs(lvel_torque[2]) * lvel_torque[2] * scl, box2 * (box0_pow4 + box1_pow4) * wp.abs(l_ang[2]) * l_ang[2] * scl,
) )
# rotate to global orientation: lfrc -> bfrc torque_global = rot @ lfrc_torque
torque = rot @ lfrc_torque force_global = rot @ lfrc_force
force = 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): def _fluid(m: Model, d: Data):
wp.launch( wp.launch(
_box_fluid, _fluid_force,
dim=(d.nworld, m.nbody), dim=(d.nworld, m.nbody),
inputs=[ inputs=[
m.opt.wind, m.opt.wind,
m.opt.density, m.opt.density,
m.opt.viscosity, m.opt.viscosity,
m.body_rootid, m.body_rootid,
m.body_geomnum,
m.body_geomadr,
m.body_mass, m.body_mass,
m.body_inertia, m.body_inertia,
m.body_fluid_ellipsoid,
m.geom_type,
m.geom_size,
m.geom_fluid,
d.xipos, d.xipos,
d.ximat, d.ximat,
d.geom_xpos,
d.geom_xmat,
d.subtree_com, d.subtree_com,
d.cvel, 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) 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), sensor_adr_to_contact_adr: wp.array(dtype=int),
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
xpos_in: wp.array2d(dtype=wp.vec3), xpos_in: wp.array2d(dtype=wp.vec3),
xipos_in: wp.array2d(dtype=wp.vec3), xipos_in: wp.array2d(dtype=wp.vec3),
geom_xpos_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( contact_forcetorque = support.contact_force_fn(
opt_cone, opt_cone,
njmax_in, njmax_in,
ncon_in, nacon_in,
contact_frame_in, contact_frame_in,
contact_friction_in, contact_friction_in,
contact_dim_in, contact_dim_in,
@@ -1896,7 +1896,7 @@ def _sensor_acc(
contact_forcetorque = support.contact_force_fn( contact_forcetorque = support.contact_force_fn(
opt_cone, opt_cone,
njmax_in, njmax_in,
ncon_in, nacon_in,
contact_frame_in, contact_frame_in,
contact_friction_in, contact_friction_in,
contact_dim_in, contact_dim_in,
@@ -1973,7 +1973,7 @@ def _sensor_acc(
contact_forcetorque = support.contact_force_fn( contact_forcetorque = support.contact_force_fn(
opt_cone, opt_cone,
njmax_in, njmax_in,
ncon_in, nacon_in,
contact_frame_in, contact_frame_in,
contact_friction_in, contact_friction_in,
contact_dim_in, contact_dim_in,
@@ -2082,7 +2082,7 @@ def _sensor_touch(
sensor_adr: wp.array(dtype=int), sensor_adr: wp.array(dtype=int),
sensor_touch_adr: wp.array(dtype=int), sensor_touch_adr: wp.array(dtype=int),
# Data in: # Data in:
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
site_xpos_in: wp.array2d(dtype=wp.vec3), site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33), site_xmat_in: wp.array2d(dtype=wp.mat33),
contact_pos_in: wp.array(dtype=wp.vec3), contact_pos_in: wp.array(dtype=wp.vec3),
@@ -2097,7 +2097,7 @@ def _sensor_touch(
): ):
conid, sensortouchadrid = wp.tid() conid, sensortouchadrid = wp.tid()
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
sensorid = sensor_touch_adr[sensortouchadrid] sensorid = sensor_touch_adr[sensortouchadrid]
@@ -2182,7 +2182,7 @@ def _sensor_tactile(
taxel_vertadr: wp.array(dtype=int), taxel_vertadr: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int), taxel_sensorid: wp.array(dtype=int),
# Data in: # Data in:
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3),
@@ -2194,7 +2194,7 @@ def _sensor_tactile(
): ):
conid, taxelid = wp.tid() conid, taxelid = wp.tid()
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
worldid = contact_worldid_in[conid] worldid = contact_worldid_in[conid]
@@ -2308,7 +2308,7 @@ def _contact_match(
sensor_contact_adr: wp.array(dtype=int), sensor_contact_adr: wp.array(dtype=int),
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
site_xpos_in: wp.array2d(dtype=wp.vec3), site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33), site_xmat_in: wp.array2d(dtype=wp.mat33),
contact_dist_in: wp.array(dtype=float), contact_dist_in: wp.array(dtype=float),
@@ -2329,7 +2329,7 @@ def _contact_match(
contactsensorid, contactid = wp.tid() contactsensorid, contactid = wp.tid()
sensorid = sensor_contact_adr[contactsensorid] sensorid = sensor_contact_adr[contactsensorid]
if contactid >= ncon_in[0]: if contactid >= nacon_in[0]:
return return
# sensor information # sensor information
@@ -2403,7 +2403,7 @@ def _contact_match(
contact_force = support.contact_force_fn( contact_force = support.contact_force_fn(
opt_cone, opt_cone,
njmax_in, njmax_in,
ncon_in, nacon_in,
contact_frame_in, contact_frame_in,
contact_friction_in, contact_friction_in,
contact_dim_in, contact_dim_in,
@@ -2422,6 +2422,7 @@ def _contact_match(
sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = dir sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = dir
@cache_kernel
def _contact_sort(maxmatch: int): def _contact_sort(maxmatch: int):
@nested_kernel(module="unique", enable_backward=False) @nested_kernel(module="unique", enable_backward=False)
def contact_sort( def contact_sort(
@@ -2466,7 +2467,7 @@ def sensor_acc(m: Model, d: Data):
wp.launch( wp.launch(
_sensor_touch, _sensor_touch,
dim=(d.nconmax, m.sensor_touch_adr.size), dim=(d.naconmax, m.sensor_touch_adr.size),
inputs=[ inputs=[
m.opt.cone, m.opt.cone,
m.geom_bodyid, m.geom_bodyid,
@@ -2476,7 +2477,7 @@ def sensor_acc(m: Model, d: Data):
m.sensor_objid, m.sensor_objid,
m.sensor_adr, m.sensor_adr,
m.sensor_touch_adr, m.sensor_touch_adr,
d.ncon, d.nacon,
d.site_xpos, d.site_xpos,
d.site_xmat, d.site_xmat,
d.contact.pos, d.contact.pos,
@@ -2494,7 +2495,7 @@ def sensor_acc(m: Model, d: Data):
wp.launch( wp.launch(
_sensor_tactile, _sensor_tactile,
dim=(d.nconmax, m.nsensortaxel), dim=(d.naconmax, m.nsensortaxel),
inputs=[ inputs=[
m.body_rootid, m.body_rootid,
m.body_weldid, m.body_weldid,
@@ -2518,7 +2519,7 @@ def sensor_acc(m: Model, d: Data):
m.geom_plugin_index, m.geom_plugin_index,
m.taxel_vertadr, m.taxel_vertadr,
m.taxel_sensorid, m.taxel_sensorid,
d.ncon, d.nacon,
d.geom_xpos, d.geom_xpos,
d.geom_xmat, d.geom_xmat,
d.subtree_com, d.subtree_com,
@@ -2539,7 +2540,7 @@ def sensor_acc(m: Model, d: Data):
wp.launch( wp.launch(
_contact_match, _contact_match,
dim=(m.sensor_contact_adr.size, d.nconmax), dim=(m.sensor_contact_adr.size, d.naconmax),
inputs=[ inputs=[
m.opt.cone, m.opt.cone,
m.opt.contact_sensor_maxmatch, m.opt.contact_sensor_maxmatch,
@@ -2554,7 +2555,7 @@ def sensor_acc(m: Model, d: Data):
m.sensor_intprm, m.sensor_intprm,
m.sensor_contact_adr, m.sensor_contact_adr,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.site_xpos, d.site_xpos,
d.site_xmat, d.site_xmat,
d.contact.dist, d.contact.dist,
@@ -2616,7 +2617,7 @@ def sensor_acc(m: Model, d: Data):
m.sensor_acc_adr, m.sensor_acc_adr,
m.sensor_adr_to_contact_adr, m.sensor_adr_to_contact_adr,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.xpos, d.xpos,
d.xipos, d.xipos,
d.geom_xpos, 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), geom_bodyid: wp.array(dtype=int),
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3),
contact_pos_in: wp.array(dtype=wp.vec3), contact_pos_in: wp.array(dtype=wp.vec3),
contact_frame_in: wp.array(dtype=wp.mat33), contact_frame_in: wp.array(dtype=wp.mat33),
@@ -1258,7 +1258,7 @@ def _cfrc_ext_contact(
): ):
contactid = wp.tid() contactid = wp.tid()
if contactid >= ncon_in[0]: if contactid >= nacon_in[0]:
return return
geom = contact_geom_in[contactid] geom = contact_geom_in[contactid]
@@ -1274,7 +1274,7 @@ def _cfrc_ext_contact(
force = support.contact_force_fn( force = support.contact_force_fn(
opt_cone, opt_cone,
njmax_in, njmax_in,
ncon_in, nacon_in,
contact_frame_in, contact_frame_in,
contact_friction_in, contact_friction_in,
contact_dim_in, contact_dim_in,
@@ -1338,13 +1338,13 @@ def rne_postconstraint(m: Model, d: Data):
# cfrc_ext += contacts # cfrc_ext += contacts
wp.launch( wp.launch(
_cfrc_ext_contact, _cfrc_ext_contact,
dim=(d.nconmax,), dim=(d.naconmax,),
inputs=[ inputs=[
m.opt.cone, m.opt.cone,
m.body_rootid, m.body_rootid,
m.geom_bodyid, m.geom_bodyid,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.subtree_com, d.subtree_com,
d.contact.pos, d.contact.pos,
d.contact.frame, d.contact.frame,
@@ -2077,7 +2077,7 @@ def _transmission_body_moment(
actuator_trnid: wp.array(dtype=wp.vec2i), actuator_trnid: wp.array(dtype=wp.vec2i),
actuator_trntype_body_adr: wp.array(dtype=int), actuator_trntype_body_adr: wp.array(dtype=int),
# Data in: # Data in:
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
subtree_com_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector), cdof_in: wp.array2d(dtype=wp.spatial_vector),
contact_dist_in: wp.array(dtype=float), contact_dist_in: wp.array(dtype=float),
@@ -2097,7 +2097,7 @@ def _transmission_body_moment(
actid = actuator_trntype_body_adr[trnbodyid] actid = actuator_trntype_body_adr[trnbodyid]
bodyid = actuator_trnid[actid][0] bodyid = actuator_trnid[actid][0]
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
worldid = contact_worldid_in[conid] worldid = contact_worldid_in[conid]
@@ -2247,11 +2247,7 @@ def transmission(m: Model, d: Data):
# compute moments # compute moments
wp.launch( wp.launch(
_transmission_body_moment, _transmission_body_moment,
dim=( dim=(m.actuator_trntype_body_adr.size, d.naconmax, m.nv),
m.actuator_trntype_body_adr.size,
d.nconmax,
m.nv,
),
inputs=[ inputs=[
m.opt.cone, m.opt.cone,
m.body_parentid, m.body_parentid,
@@ -2260,7 +2256,7 @@ def transmission(m: Model, d: Data):
m.geom_bodyid, m.geom_bodyid,
m.actuator_trnid, m.actuator_trnid,
m.actuator_trntype_body_adr, m.actuator_trntype_body_adr,
d.ncon, d.nacon,
d.subtree_com, d.subtree_com,
d.cdof, d.cdof,
d.contact.dist, 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, opt_ls_parallel_min_step: float,
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
ne_in: wp.array(dtype=int), ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int), nf_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int), nefc_in: wp.array(dtype=int),
@@ -561,7 +561,7 @@ def linesearch_parallel_fused(
# extract contact info # extract contact info
conid = efc_id_in[worldid, efcid] conid = efc_id_in[worldid, efcid]
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
continue continue
efcid0 = contact_efc_address_in[conid, 0] 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.impratio,
m.opt.ls_parallel_min_step, m.opt.ls_parallel_min_step,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.ne, d.ne,
d.nf, d.nf,
d.nefc, d.nefc,
@@ -771,7 +771,7 @@ def linesearch_prepare_quad(
# Model: # Model:
opt_impratio: wp.array(dtype=float), opt_impratio: wp.array(dtype=float),
# Data in: # Data in:
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int), nefc_in: wp.array(dtype=int),
contact_friction_in: wp.array(dtype=types.vec5), contact_friction_in: wp.array(dtype=types.vec5),
contact_dim_in: wp.array(dtype=int), contact_dim_in: wp.array(dtype=int),
@@ -805,7 +805,7 @@ def linesearch_prepare_quad(
# extract contact info # extract contact info
conid = efc_id_in[worldid, efcid] conid = efc_id_in[worldid, efcid]
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
efcid0 = contact_efc_address_in[conid, 0] efcid0 = contact_efc_address_in[conid, 0]
@@ -951,7 +951,7 @@ def _linesearch(m: types.Model, d: types.Data):
dim=(d.nworld, d.njmax), dim=(d.nworld, d.njmax),
inputs=[ inputs=[
m.opt.impratio, m.opt.impratio,
d.ncon, d.nacon,
d.nefc, d.nefc,
d.contact.friction, d.contact.friction,
d.contact.dim, d.contact.dim,
@@ -1064,7 +1064,7 @@ def update_constraint_efc(
# Model: # Model:
opt_impratio: wp.array(dtype=float), opt_impratio: wp.array(dtype=float),
# Data in: # Data in:
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
ne_in: wp.array(dtype=int), ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int), nf_in: wp.array(dtype=int),
nefc_in: wp.array(dtype=int), nefc_in: wp.array(dtype=int),
@@ -1129,7 +1129,7 @@ def update_constraint_efc(
else: # elliptic friction cone contact else: # elliptic friction cone contact
conid = efc_id_in[worldid, efcid] conid = efc_id_in[worldid, efcid]
if conid >= ncon_in[0]: if conid >= nacon_in[0]:
return return
dim = contact_dim_in[conid] dim = contact_dim_in[conid]
@@ -1264,7 +1264,7 @@ def _update_constraint(m: types.Model, d: types.Data):
dim=(d.nworld, d.njmax), dim=(d.nworld, d.njmax),
inputs=[ inputs=[
m.opt.impratio, m.opt.impratio,
d.ncon, d.nacon,
d.ne, d.ne,
d.nf, d.nf,
d.nefc, d.nefc,
@@ -1468,8 +1468,8 @@ def update_gradient_JTCJ(
dof_tri_row: wp.array(dtype=int), dof_tri_row: wp.array(dtype=int),
dof_tri_col: wp.array(dtype=int), dof_tri_col: wp.array(dtype=int),
# Data in: # Data in:
nconmax_in: int, naconmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
contact_dist_in: wp.array(dtype=float), contact_dist_in: wp.array(dtype=float),
contact_includemargin_in: wp.array(dtype=float), contact_includemargin_in: wp.array(dtype=float),
contact_friction_in: wp.array(dtype=types.vec5), contact_friction_in: wp.array(dtype=types.vec5),
@@ -1495,7 +1495,7 @@ def update_gradient_JTCJ(
for i in range(nblocks_perblock): for i in range(nblocks_perblock):
conid = conid_start + i * dim_block conid = conid_start + i * dim_block
if conid >= min(ncon_in[0], nconmax_in): if conid >= min(nacon_in[0], naconmax_in):
return return
worldid = contact_worldid_in[conid] 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) dim_block = ceil((sm_count * 6 * 256) / m.dof_tri_row.size)
else: else:
# fall back for CPU # 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( wp.launch(
update_gradient_JTCJ, update_gradient_JTCJ,
@@ -1738,8 +1738,8 @@ def _update_gradient(m: types.Model, d: types.Data):
m.opt.impratio, m.opt.impratio,
m.dof_tri_row, m.dof_tri_row,
m.dof_tri_col, m.dof_tri_col,
d.nconmax, d.naconmax,
d.ncon, d.nacon,
d.contact.dist, d.contact.dist,
d.contact.includemargin, d.contact.includemargin,
d.contact.friction, 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, opt_cone: int,
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
contact_frame_in: wp.array(dtype=wp.mat33), contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5), contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int), contact_dim_in: wp.array(dtype=int),
@@ -312,7 +312,7 @@ def contact_force_fn(
condim = contact_dim_in[contact_id] condim = contact_dim_in[contact_id]
efc_address = contact_efc_address_in[contact_id, 0] 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: if opt_cone == ConeType.PYRAMIDAL:
force = _decode_pyramid( force = _decode_pyramid(
njmax_in, njmax_in,
@@ -341,7 +341,7 @@ def contact_force_kernel(
opt_cone: int, opt_cone: int,
# Data in: # Data in:
njmax_in: int, njmax_in: int,
ncon_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int),
contact_frame_in: wp.array(dtype=wp.mat33), contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5), contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int), contact_dim_in: wp.array(dtype=int),
@@ -358,7 +358,7 @@ def contact_force_kernel(
contactid = contact_ids[tid] contactid = contact_ids[tid]
if contactid >= ncon_in[0]: if contactid >= nacon_in[0]:
return return
worldid = contact_worldid_in[contactid] worldid = contact_worldid_in[contactid]
@@ -366,7 +366,7 @@ def contact_force_kernel(
out[tid] = contact_force_fn( out[tid] = contact_force_fn(
opt_cone, opt_cone,
njmax_in, njmax_in,
ncon_in, nacon_in,
contact_frame_in, contact_frame_in,
contact_friction_in, contact_friction_in,
contact_dim_in, contact_dim_in,
@@ -401,7 +401,7 @@ def contact_force(
inputs=[ inputs=[
m.opt.cone, m.opt.cone,
d.njmax, d.njmax,
d.ncon, d.nacon,
d.contact.frame, d.contact.frame,
d.contact.friction, d.contact.friction,
d.contact.dim, 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_contype: OR over all geom contypes (nbody,)
body_conaffinity: OR over all geom conaffinities (nbody,) body_conaffinity: OR over all geom conaffinities (nbody,)
body_gravcomp: antigravity force, units of body weight (nworld, 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_type: type of joint (JointType) (njnt,)
jnt_qposadr: start addr in 'qpos' for joint's data (njnt,) jnt_qposadr: start addr in 'qpos' for joint's data (njnt,)
jnt_dofadr: start addr in 'qvel' 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_solref: constraint solver reference: contact (nworld, ngeom, mjNREF)
geom_solimp: constraint solver impedance: contact (nworld, ngeom, mjNIMP) geom_solimp: constraint solver impedance: contact (nworld, ngeom, mjNIMP)
geom_size: geom-specific size parameters (ngeom, 3) geom_size: geom-specific size parameters (ngeom, 3)
geom_fluid: fluid interaction parameters (ngeom, mjNFLUID)
geom_aabb: bounding box, (center, size) (ngeom, 6) geom_aabb: bounding box, (center, size) (ngeom, 6)
geom_rbound: radius of bounding sphere (nworld, ngeom,) geom_rbound: radius of bounding sphere (nworld, ngeom,)
geom_pos: local position offset rel. to body (nworld, ngeom, 3) geom_pos: local position offset rel. to body (nworld, ngeom, 3)
@@ -1116,6 +1118,7 @@ class Model:
body_contype: wp.array(dtype=int) body_contype: wp.array(dtype=int)
body_conaffinity: wp.array(dtype=int) body_conaffinity: wp.array(dtype=int)
body_gravcomp: wp.array2d(dtype=float) body_gravcomp: wp.array2d(dtype=float)
body_fluid_ellipsoid: wp.array(dtype=bool)
jnt_type: wp.array(dtype=int) jnt_type: wp.array(dtype=int)
jnt_qposadr: wp.array(dtype=int) jnt_qposadr: wp.array(dtype=int)
jnt_dofadr: wp.array(dtype=int) jnt_dofadr: wp.array(dtype=int)
@@ -1158,6 +1161,7 @@ class Model:
geom_solref: wp.array2d(dtype=wp.vec2) geom_solref: wp.array2d(dtype=wp.vec2)
geom_solimp: wp.array2d(dtype=vec5) geom_solimp: wp.array2d(dtype=vec5)
geom_size: wp.array2d(dtype=wp.vec3) geom_size: wp.array2d(dtype=wp.vec3)
geom_fluid: wp.array2d(dtype=float)
geom_aabb: wp.array2d(dtype=wp.vec3) geom_aabb: wp.array2d(dtype=wp.vec3)
geom_rbound: wp.array2d(dtype=float) geom_rbound: wp.array2d(dtype=float)
geom_pos: wp.array2d(dtype=wp.vec3) geom_pos: wp.array2d(dtype=wp.vec3)
@@ -1365,18 +1369,18 @@ class Contact:
"""Contact data. """Contact data.
Attributes: Attributes:
dist: distance between nearest points; neg: penetration dist: distance between nearest points; neg: penetration (naconmax,)
pos: position of contact point: midpoint between geoms pos: position of contact point: midpoint between geoms (naconmax, 3)
frame: normal is in [0-2], points from geom[0] to geom[1] frame: normal is in [0-2], points from geom[0] to geom[1] (naconmax, 3, 3)
includemargin: include if dist<includemargin=margin-gap includemargin: include if dist<includemargin=margin-gap (naconmax,)
friction: tangent1, 2, spin, roll1, 2 friction: tangent1, 2, spin, roll1, 2 (naconmax, 5)
solref: constraint solver reference, normal direction solref: constraint solver reference, normal direction (naconmax, 2)
solreffriction: constraint solver reference, friction directions solreffriction: constraint solver reference, friction directions (naconmax, 2)
solimp: constraint solver impedance solimp: constraint solver impedance (naconmax, 5)
dim: contact space dimensionality: 1, 3, 4 or 6 dim: contact space dimensionality: 1, 3, 4 or 6 (naconmax,)
geom: geom ids; -1 for flex geom: geom ids; -1 for flex (naconmax, 2)
efc_address: address in efc; -1: not included efc_address: address in efc; -1: not included (naconmax, ncondim)
worldid: world id worldid: world id (naconmax,)
""" """
dist: wp.array(dtype=float) dist: wp.array(dtype=float)
@@ -1399,10 +1403,10 @@ class Data:
Attributes: Attributes:
nworld: number of worlds 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 njmax: maximum number of constraints per world
solver_niter: number of solver iterations (nworld,) 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: number of equality constraints (nworld,)
ne_connect: number of equality connect constraints (nworld,) ne_connect: number of equality connect constraints (nworld,)
ne_weld: number of equality weld constraints (nworld,) ne_weld: number of equality weld constraints (nworld,)
@@ -1498,31 +1502,31 @@ class Data:
sap_cumulative_sum: broadphase context (nworld, ngeom) sap_cumulative_sum: broadphase context (nworld, ngeom)
sap_segment_index: broadphase context (requires nworld + 1) (nworld, 2) sap_segment_index: broadphase context (requires nworld + 1) (nworld, 2)
dyn_geom_aabb: dynamic geometry axis-aligned bounding boxes (nworld, ngeom, 2) dyn_geom_aabb: dynamic geometry axis-aligned bounding boxes (nworld, ngeom, 2)
collision_pair: collision pairs from broadphase (nconmax,) collision_pair: collision pairs from broadphase (naconmax,)
collision_worldid: collision world ids from broadphase (nconmax,) collision_worldid: collision world ids from broadphase (naconmax,)
ncollision: collision count from broadphase ncollision: collision count from broadphase
epa_vert: vertices in EPA polytope in Minkowski space (nconmax, 5 + CCDiter) epa_vert: vertices in EPA polytope in Minkowski space (naconmax, 5 + CCDiter)
epa_vert1: vertices in EPA polytope in geom 1 space (nconmax, 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 (nconmax, 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 (nconmax, 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 (nconmax, 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 (nconmax, 6 + 5 * CCDiter) epa_face: faces of polytope represented by three indices (naconmax, 6 + 5 * CCDiter)
epa_pr: projection of origin on polytope faces (nconmax, 6 + 5 * CCDiter) epa_pr: projection of origin on polytope faces (naconmax, 6 + 5 * CCDiter)
epa_norm2: epa_pr * epa_pr (nconmax, 6 + 5 * CCDiter) epa_norm2: epa_pr * epa_pr (naconmax, 6 + 5 * CCDiter)
epa_index: index of face in polytope map (nconmax, 6 + 5 * CCDiter) epa_index: index of face in polytope map (naconmax, 6 + 5 * CCDiter)
epa_map: status of faces in polytope (nconmax, 6 + 5 * CCDiter) epa_map: status of faces in polytope (naconmax, 6 + 5 * CCDiter)
epa_horizon: index pair (i j) of edges on horizon (nconmax, 2 * 12) epa_horizon: index pair (i j) of edges on horizon (naconmax, 2 * 12)
multiccd_polygon: clipped contact surface (nconmax, 2 * max_npolygon) multiccd_polygon: clipped contact surface (naconmax, 2 * max_npolygon)
multiccd_clipped: clipped contact surface (intermediate) (nconmax, 2 * max_npolygon) multiccd_clipped: clipped contact surface (intermediate) (naconmax, 2 * max_npolygon)
multiccd_pnormal: plane normal of clipping polygon (nconmax, max_npolygon) multiccd_pnormal: plane normal of clipping polygon (naconmax, max_npolygon)
multiccd_pdist: plane distance of clipping polygon (nconmax, max_npolygon) multiccd_pdist: plane distance of clipping polygon (naconmax, max_npolygon)
multiccd_idx1: list of normal index candidates for Geom 1 (nconmax, max_meshdegree) multiccd_idx1: list of normal index candidates for Geom 1 (naconmax, max_meshdegree)
multiccd_idx2: list of normal index candidates for Geom 2 (nconmax, max_meshdegree) multiccd_idx2: list of normal index candidates for Geom 2 (naconmax, max_meshdegree)
multiccd_n1: list of normal candidates for Geom 1 (nconmax, max_meshdegree) multiccd_n1: list of normal candidates for Geom 1 (naconmax, max_meshdegree)
multiccd_n2: list of normal candidates for Geom 1 (nconmax, max_meshdegree) multiccd_n2: list of normal candidates for Geom 1 (naconmax, max_meshdegree)
multiccd_endvert: list of edge vertices candidates (nconmax, max_meshdegree) multiccd_endvert: list of edge vertices candidates (naconmax, max_meshdegree)
multiccd_face1: contact face (nconmax, max_npolygon) multiccd_face1: contact face (naconmax, max_npolygon)
multiccd_face2: contact face (nconmax, max_npolygon) multiccd_face2: contact face (naconmax, max_npolygon)
cacc: com-based acceleration (nworld, nbody, 6) cacc: com-based acceleration (nworld, nbody, 6)
cfrc_int: com-based interaction force with parent (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) cfrc_ext: com-based external force on body (nworld, nbody, 6)
@@ -1554,10 +1558,10 @@ class Data:
""" """
nworld: int # warp only nworld: int # warp only
nconmax: int # warp only naconmax: int # warp only
njmax: int # warp only njmax: int # warp only
solver_niter: wp.array(dtype=int) solver_niter: wp.array(dtype=int)
ncon: wp.array(dtype=int) nacon: wp.array(dtype=int)
ne: wp.array(dtype=int) ne: wp.array(dtype=int)
ne_connect: wp.array(dtype=int) # warp only ne_connect: wp.array(dtype=int) # warp only
ne_weld: 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() tid = wp.tid()
b[tid] = a[tid] + 1.0 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: Args:
f: The function to be registered as a kernel. f: The function to be registered as a kernel.
enable_backward: If False, the backward pass will not be generated. enable_backward: If False, the backward pass will not be generated.
@@ -168,15 +175,27 @@ def kernel(
Returns: Returns:
The registered kernel. 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 = {} _KERNEL_CACHE = {}
@@ -186,7 +205,14 @@ def cache_kernel(func):
# caching kernels to avoid crashes in graph_conditional code # caching kernels to avoid crashes in graph_conditional code
@functools.wraps(func) @functools.wraps(func)
def wrapper(*args): 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: if key not in _KERNEL_CACHE:
_KERNEL_CACHE[key] = func(*args) _KERNEL_CACHE[key] = func(*args)
return _KERNEL_CACHE[key] 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: def _load_model(path: epath.Path) -> mujoco.MjModel:
if not path.exists(): 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(): if not resource_path.exists():
raise FileNotFoundError(f"file not found: {path}\nalso tried: {resource_path}") raise FileNotFoundError(f"file not found: {path}\nalso tried: {resource_path}")
path = 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_initpoints: int,
opt__sdf_iterations: int, opt__sdf_iterations: int,
# Data # Data
nconmax: int, naconmax: int,
collision_pair: wp.array(dtype=wp.vec2i), collision_pair: wp.array(dtype=wp.vec2i),
collision_pairid: wp.array(dtype=int), collision_pairid: wp.array(dtype=int),
collision_worldid: wp.array(dtype=int), collision_worldid: wp.array(dtype=int),
@@ -140,8 +140,8 @@ def _collision_shim(
multiccd_pdist: wp.array2d(dtype=float), multiccd_pdist: wp.array2d(dtype=float),
multiccd_pnormal: wp.array2d(dtype=wp.vec3), multiccd_pnormal: wp.array2d(dtype=wp.vec3),
multiccd_polygon: wp.array2d(dtype=wp.vec3), multiccd_polygon: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int), ncollision: wp.array(dtype=int),
ncon: wp.array(dtype=int),
sap_cumulative_sum: wp.array2d(dtype=int), sap_cumulative_sum: wp.array2d(dtype=int),
sap_projection_lower: wp.array3d(dtype=float), sap_projection_lower: wp.array3d(dtype=float),
sap_projection_upper: wp.array2d(dtype=float), sap_projection_upper: wp.array2d(dtype=float),
@@ -266,9 +266,9 @@ def _collision_shim(
_d.multiccd_pdist = multiccd_pdist _d.multiccd_pdist = multiccd_pdist
_d.multiccd_pnormal = multiccd_pnormal _d.multiccd_pnormal = multiccd_pnormal
_d.multiccd_polygon = multiccd_polygon _d.multiccd_polygon = multiccd_polygon
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision _d.ncollision = ncollision
_d.ncon = ncon
_d.nconmax = nconmax
_d.sap_cumulative_sum = sap_cumulative_sum _d.sap_cumulative_sum = sap_cumulative_sum
_d.sap_projection_lower = sap_projection_lower _d.sap_projection_lower = sap_projection_lower
_d.sap_projection_upper = sap_projection_upper _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_pdist': d._impl.multiccd_pdist.shape,
'multiccd_pnormal': d._impl.multiccd_pnormal.shape, 'multiccd_pnormal': d._impl.multiccd_pnormal.shape,
'multiccd_polygon': d._impl.multiccd_polygon.shape, 'multiccd_polygon': d._impl.multiccd_polygon.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape, 'ncollision': d._impl.ncollision.shape,
'ncon': d._impl.ncon.shape,
'sap_cumulative_sum': d._impl.sap_cumulative_sum.shape, 'sap_cumulative_sum': d._impl.sap_cumulative_sum.shape,
'sap_projection_lower': d._impl.sap_projection_lower.shape, 'sap_projection_lower': d._impl.sap_projection_lower.shape,
'sap_projection_upper': d._impl.sap_projection_upper.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_pdist',
'multiccd_pnormal', 'multiccd_pnormal',
'multiccd_polygon', 'multiccd_polygon',
'nacon',
'ncollision', 'ncollision',
'ncon',
'sap_cumulative_sum', 'sap_cumulative_sum',
'sap_projection_lower', 'sap_projection_lower',
'sap_projection_upper', '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.legacy_gjk,
m.opt._impl.sdf_initpoints, m.opt._impl.sdf_initpoints,
m.opt._impl.sdf_iterations, m.opt._impl.sdf_iterations,
d._impl.nconmax, d._impl.naconmax,
d._impl.collision_pair, d._impl.collision_pair,
d._impl.collision_pairid, d._impl.collision_pairid,
d._impl.collision_worldid, 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_pdist,
d._impl.multiccd_pnormal, d._impl.multiccd_pnormal,
d._impl.multiccd_polygon, d._impl.multiccd_polygon,
d._impl.nacon,
d._impl.ncollision, d._impl.ncollision,
d._impl.ncon,
d._impl.sap_cumulative_sum, d._impl.sap_cumulative_sum,
d._impl.sap_projection_lower, d._impl.sap_projection_lower,
d._impl.sap_projection_upper, 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_pdist': out[24],
'_impl.multiccd_pnormal': out[25], '_impl.multiccd_pnormal': out[25],
'_impl.multiccd_polygon': out[26], '_impl.multiccd_polygon': out[26],
'_impl.ncollision': out[27], '_impl.nacon': out[27],
'_impl.ncon': out[28], '_impl.ncollision': out[28],
'_impl.sap_cumulative_sum': out[29], '_impl.sap_cumulative_sum': out[29],
'_impl.sap_projection_lower': out[30], '_impl.sap_projection_lower': out[30],
'_impl.sap_projection_upper': out[31], '_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 @ffi.marshal_jax_warp_callable
def collision(m: types.Model, d: types.Data): def collision(m: types.Model, d: types.Data):
return _collision_jax_impl(m, d) return _collision_jax_impl(m, d)
@collision.def_vmap @collision.def_vmap
@ffi.marshal_custom_vmap @ffi.marshal_custom_vmap
def collision_vmap(unused_axis_size, is_batched, m, d): 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. # Check leading dimensions.
jax.tree.map_with_path( jax.tree.map_with_path(
lambda path, x: _check_leading_dim( 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, d_broadcast,
) )
mjx/mujoco/mjx/warp/forward.py
+32 -24
View File
@@ -74,6 +74,7 @@ def _forward_shim(
block_dim: mjwp_types.BlockDim, block_dim: mjwp_types.BlockDim,
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array(dtype=int),
body_fluid_ellipsoid: wp.array(dtype=bool),
body_geomadr: wp.array(dtype=int), body_geomadr: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int), body_geomnum: wp.array(dtype=int),
body_gravcomp: wp.array2d(dtype=float), body_gravcomp: wp.array2d(dtype=float),
@@ -144,6 +145,7 @@ def _forward_shim(
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int),
geom_condim: wp.array(dtype=int), geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array(dtype=int),
geom_fluid: wp.array2d(dtype=float),
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3),
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d(dtype=float),
geom_group: wp.array(dtype=int), geom_group: wp.array(dtype=int),
@@ -351,7 +353,7 @@ def _forward_shim(
opt__wind: wp.array(dtype=wp.vec3), opt__wind: wp.array(dtype=wp.vec3),
stat__meaninertia: float, stat__meaninertia: float,
# Data # Data
nconmax: int, naconmax: int,
njmax: int, njmax: int,
act: wp.array2d(dtype=float), act: wp.array2d(dtype=float),
act_dot: wp.array2d(dtype=float), act_dot: wp.array2d(dtype=float),
@@ -409,8 +411,8 @@ def _forward_shim(
multiccd_pdist: wp.array2d(dtype=float), multiccd_pdist: wp.array2d(dtype=float),
multiccd_pnormal: wp.array2d(dtype=wp.vec3), multiccd_pnormal: wp.array2d(dtype=wp.vec3),
multiccd_polygon: wp.array2d(dtype=wp.vec3), multiccd_polygon: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int), ncollision: wp.array(dtype=int),
ncon: wp.array(dtype=int),
ne: wp.array(dtype=int), ne: wp.array(dtype=int),
ne_connect: wp.array(dtype=int), ne_connect: wp.array(dtype=int),
ne_jnt: wp.array(dtype=int), ne_jnt: wp.array(dtype=int),
@@ -558,6 +560,7 @@ def _forward_shim(
_m.block_dim = block_dim _m.block_dim = block_dim
_m.body_dofadr = body_dofadr _m.body_dofadr = body_dofadr
_m.body_dofnum = body_dofnum _m.body_dofnum = body_dofnum
_m.body_fluid_ellipsoid = body_fluid_ellipsoid
_m.body_geomadr = body_geomadr _m.body_geomadr = body_geomadr
_m.body_geomnum = body_geomnum _m.body_geomnum = body_geomnum
_m.body_gravcomp = body_gravcomp _m.body_gravcomp = body_gravcomp
@@ -628,6 +631,7 @@ def _forward_shim(
_m.geom_bodyid = geom_bodyid _m.geom_bodyid = geom_bodyid
_m.geom_condim = geom_condim _m.geom_condim = geom_condim
_m.geom_dataid = geom_dataid _m.geom_dataid = geom_dataid
_m.geom_fluid = geom_fluid
_m.geom_friction = geom_friction _m.geom_friction = geom_friction
_m.geom_gap = geom_gap _m.geom_gap = geom_gap
_m.geom_group = geom_group _m.geom_group = geom_group
@@ -936,9 +940,9 @@ def _forward_shim(
_d.multiccd_pdist = multiccd_pdist _d.multiccd_pdist = multiccd_pdist
_d.multiccd_pnormal = multiccd_pnormal _d.multiccd_pnormal = multiccd_pnormal
_d.multiccd_polygon = multiccd_polygon _d.multiccd_polygon = multiccd_polygon
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision _d.ncollision = ncollision
_d.ncon = ncon
_d.nconmax = nconmax
_d.ne = ne _d.ne = ne
_d.ne_connect = ne_connect _d.ne_connect = ne_connect
_d.ne_jnt = ne_jnt _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_pdist': d._impl.multiccd_pdist.shape,
'multiccd_pnormal': d._impl.multiccd_pnormal.shape, 'multiccd_pnormal': d._impl.multiccd_pnormal.shape,
'multiccd_polygon': d._impl.multiccd_polygon.shape, 'multiccd_polygon': d._impl.multiccd_polygon.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape, 'ncollision': d._impl.ncollision.shape,
'ncon': d._impl.ncon.shape,
'ne': d._impl.ne.shape, 'ne': d._impl.ne.shape,
'ne_connect': d._impl.ne_connect.shape, 'ne_connect': d._impl.ne_connect.shape,
'ne_jnt': d._impl.ne_jnt.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_pdist',
'multiccd_pnormal', 'multiccd_pnormal',
'multiccd_polygon', 'multiccd_polygon',
'nacon',
'ncollision', 'ncollision',
'ncon',
'ne', 'ne',
'ne_connect', 'ne_connect',
'ne_jnt', 'ne_jnt',
@@ -1399,6 +1403,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.block_dim, m._impl.block_dim,
m.body_dofadr, m.body_dofadr,
m.body_dofnum, m.body_dofnum,
m._impl.body_fluid_ellipsoid,
m.body_geomadr, m.body_geomadr,
m.body_geomnum, m.body_geomnum,
m.body_gravcomp, m.body_gravcomp,
@@ -1469,6 +1474,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.geom_bodyid, m.geom_bodyid,
m.geom_condim, m.geom_condim,
m.geom_dataid, m.geom_dataid,
m.geom_fluid,
m.geom_friction, m.geom_friction,
m.geom_gap, m.geom_gap,
m.geom_group, m.geom_group,
@@ -1675,7 +1681,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.opt.viscosity, m.opt.viscosity,
m.opt.wind, m.opt.wind,
m.stat.meaninertia, m.stat.meaninertia,
d._impl.nconmax, d._impl.naconmax,
d._impl.njmax, d._impl.njmax,
d.act, d.act,
d.act_dot, d.act_dot,
@@ -1733,8 +1739,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.multiccd_pdist, d._impl.multiccd_pdist,
d._impl.multiccd_pnormal, d._impl.multiccd_pnormal,
d._impl.multiccd_polygon, d._impl.multiccd_polygon,
d._impl.nacon,
d._impl.ncollision, d._impl.ncollision,
d._impl.ncon,
d._impl.ne, d._impl.ne,
d._impl.ne_connect, d._impl.ne_connect,
d._impl.ne_jnt, 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_pdist': out[53],
'_impl.multiccd_pnormal': out[54], '_impl.multiccd_pnormal': out[54],
'_impl.multiccd_polygon': out[55], '_impl.multiccd_polygon': out[55],
'_impl.ncollision': out[56], '_impl.nacon': out[56],
'_impl.ncon': out[57], '_impl.ncollision': out[57],
'_impl.ne': out[58], '_impl.ne': out[58],
'_impl.ne_connect': out[59], '_impl.ne_connect': out[59],
'_impl.ne_jnt': out[60], '_impl.ne_jnt': out[60],
@@ -2033,8 +2039,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable @ffi.marshal_jax_warp_callable
def forward(m: types.Model, d: types.Data): def forward(m: types.Model, d: types.Data):
return _forward_jax_impl(m, d) return _forward_jax_impl(m, d)
@forward.def_vmap @forward.def_vmap
@ffi.marshal_custom_vmap @ffi.marshal_custom_vmap
def forward_vmap(unused_axis_size, is_batched, m, d): def forward_vmap(unused_axis_size, is_batched, m, d):
@@ -2093,6 +2097,7 @@ def _step_shim(
block_dim: mjwp_types.BlockDim, block_dim: mjwp_types.BlockDim,
body_dofadr: wp.array(dtype=int), body_dofadr: wp.array(dtype=int),
body_dofnum: wp.array(dtype=int), body_dofnum: wp.array(dtype=int),
body_fluid_ellipsoid: wp.array(dtype=bool),
body_geomadr: wp.array(dtype=int), body_geomadr: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int), body_geomnum: wp.array(dtype=int),
body_gravcomp: wp.array2d(dtype=float), body_gravcomp: wp.array2d(dtype=float),
@@ -2163,6 +2168,7 @@ def _step_shim(
geom_bodyid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int),
geom_condim: wp.array(dtype=int), geom_condim: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int), geom_dataid: wp.array(dtype=int),
geom_fluid: wp.array2d(dtype=float),
geom_friction: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3),
geom_gap: wp.array2d(dtype=float), geom_gap: wp.array2d(dtype=float),
geom_group: wp.array(dtype=int), geom_group: wp.array(dtype=int),
@@ -2371,7 +2377,7 @@ def _step_shim(
opt__wind: wp.array(dtype=wp.vec3), opt__wind: wp.array(dtype=wp.vec3),
stat__meaninertia: float, stat__meaninertia: float,
# Data # Data
nconmax: int, naconmax: int,
njmax: int, njmax: int,
act: wp.array2d(dtype=float), act: wp.array2d(dtype=float),
act_dot: wp.array2d(dtype=float), act_dot: wp.array2d(dtype=float),
@@ -2432,8 +2438,8 @@ def _step_shim(
multiccd_pdist: wp.array2d(dtype=float), multiccd_pdist: wp.array2d(dtype=float),
multiccd_pnormal: wp.array2d(dtype=wp.vec3), multiccd_pnormal: wp.array2d(dtype=wp.vec3),
multiccd_polygon: wp.array2d(dtype=wp.vec3), multiccd_polygon: wp.array2d(dtype=wp.vec3),
nacon: wp.array(dtype=int),
ncollision: wp.array(dtype=int), ncollision: wp.array(dtype=int),
ncon: wp.array(dtype=int),
ne: wp.array(dtype=int), ne: wp.array(dtype=int),
ne_connect: wp.array(dtype=int), ne_connect: wp.array(dtype=int),
ne_jnt: wp.array(dtype=int), ne_jnt: wp.array(dtype=int),
@@ -2590,6 +2596,7 @@ def _step_shim(
_m.block_dim = block_dim _m.block_dim = block_dim
_m.body_dofadr = body_dofadr _m.body_dofadr = body_dofadr
_m.body_dofnum = body_dofnum _m.body_dofnum = body_dofnum
_m.body_fluid_ellipsoid = body_fluid_ellipsoid
_m.body_geomadr = body_geomadr _m.body_geomadr = body_geomadr
_m.body_geomnum = body_geomnum _m.body_geomnum = body_geomnum
_m.body_gravcomp = body_gravcomp _m.body_gravcomp = body_gravcomp
@@ -2660,6 +2667,7 @@ def _step_shim(
_m.geom_bodyid = geom_bodyid _m.geom_bodyid = geom_bodyid
_m.geom_condim = geom_condim _m.geom_condim = geom_condim
_m.geom_dataid = geom_dataid _m.geom_dataid = geom_dataid
_m.geom_fluid = geom_fluid
_m.geom_friction = geom_friction _m.geom_friction = geom_friction
_m.geom_gap = geom_gap _m.geom_gap = geom_gap
_m.geom_group = geom_group _m.geom_group = geom_group
@@ -2972,9 +2980,9 @@ def _step_shim(
_d.multiccd_pdist = multiccd_pdist _d.multiccd_pdist = multiccd_pdist
_d.multiccd_pnormal = multiccd_pnormal _d.multiccd_pnormal = multiccd_pnormal
_d.multiccd_polygon = multiccd_polygon _d.multiccd_polygon = multiccd_polygon
_d.nacon = nacon
_d.naconmax = naconmax
_d.ncollision = ncollision _d.ncollision = ncollision
_d.ncon = ncon
_d.nconmax = nconmax
_d.ne = ne _d.ne = ne
_d.ne_connect = ne_connect _d.ne_connect = ne_connect
_d.ne_jnt = ne_jnt _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_pdist': d._impl.multiccd_pdist.shape,
'multiccd_pnormal': d._impl.multiccd_pnormal.shape, 'multiccd_pnormal': d._impl.multiccd_pnormal.shape,
'multiccd_polygon': d._impl.multiccd_polygon.shape, 'multiccd_polygon': d._impl.multiccd_polygon.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape, 'ncollision': d._impl.ncollision.shape,
'ncon': d._impl.ncon.shape,
'ne': d._impl.ne.shape, 'ne': d._impl.ne.shape,
'ne_connect': d._impl.ne_connect.shape, 'ne_connect': d._impl.ne_connect.shape,
'ne_jnt': d._impl.ne_jnt.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_pdist',
'multiccd_pnormal', 'multiccd_pnormal',
'multiccd_polygon', 'multiccd_polygon',
'nacon',
'ncollision', 'ncollision',
'ncon',
'ne', 'ne',
'ne_connect', 'ne_connect',
'ne_jnt', 'ne_jnt',
@@ -3468,6 +3476,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.block_dim, m._impl.block_dim,
m.body_dofadr, m.body_dofadr,
m.body_dofnum, m.body_dofnum,
m._impl.body_fluid_ellipsoid,
m.body_geomadr, m.body_geomadr,
m.body_geomnum, m.body_geomnum,
m.body_gravcomp, m.body_gravcomp,
@@ -3538,6 +3547,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.geom_bodyid, m.geom_bodyid,
m.geom_condim, m.geom_condim,
m.geom_dataid, m.geom_dataid,
m.geom_fluid,
m.geom_friction, m.geom_friction,
m.geom_gap, m.geom_gap,
m.geom_group, m.geom_group,
@@ -3745,7 +3755,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.opt.viscosity, m.opt.viscosity,
m.opt.wind, m.opt.wind,
m.stat.meaninertia, m.stat.meaninertia,
d._impl.nconmax, d._impl.naconmax,
d._impl.njmax, d._impl.njmax,
d.act, d.act,
d.act_dot, d.act_dot,
@@ -3806,8 +3816,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.multiccd_pdist, d._impl.multiccd_pdist,
d._impl.multiccd_pnormal, d._impl.multiccd_pnormal,
d._impl.multiccd_polygon, d._impl.multiccd_polygon,
d._impl.nacon,
d._impl.ncollision, d._impl.ncollision,
d._impl.ncon,
d._impl.ne, d._impl.ne,
d._impl.ne_connect, d._impl.ne_connect,
d._impl.ne_jnt, 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_pdist': out[56],
'_impl.multiccd_pnormal': out[57], '_impl.multiccd_pnormal': out[57],
'_impl.multiccd_polygon': out[58], '_impl.multiccd_polygon': out[58],
'_impl.ncollision': out[59], '_impl.nacon': out[59],
'_impl.ncon': out[60], '_impl.ncollision': out[60],
'_impl.ne': out[61], '_impl.ne': out[61],
'_impl.ne_connect': out[62], '_impl.ne_connect': out[62],
'_impl.ne_jnt': out[63], '_impl.ne_jnt': out[63],
@@ -4127,8 +4137,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable @ffi.marshal_jax_warp_callable
def step(m: types.Model, d: types.Data): def step(m: types.Model, d: types.Data):
return _step_jax_impl(m, d) return _step_jax_impl(m, d)
@step.def_vmap @step.def_vmap
@ffi.marshal_custom_vmap @ffi.marshal_custom_vmap
def step_vmap(unused_axis_size, is_batched, m, d): 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} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _kinematics_shim( def _kinematics_shim(
# Model # Model
@@ -281,8 +280,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable @ffi.marshal_jax_warp_callable
def kinematics(m: types.Model, d: types.Data): def kinematics(m: types.Model, d: types.Data):
return _kinematics_jax_impl(m, d) return _kinematics_jax_impl(m, d)
@kinematics.def_vmap @kinematics.def_vmap
@ffi.marshal_custom_vmap @ffi.marshal_custom_vmap
def kinematics_vmap(unused_axis_size, is_batched, m, d): 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} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _tendon_shim( def _tendon_shim(
# Model # Model
@@ -491,8 +487,6 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
@ffi.marshal_jax_warp_callable @ffi.marshal_jax_warp_callable
def tendon(m: types.Model, d: types.Data): def tendon(m: types.Model, d: types.Data):
return _tendon_jax_impl(m, d) return _tendon_jax_impl(m, d)
@tendon.def_vmap @tendon.def_vmap
@ffi.marshal_custom_vmap @ffi.marshal_custom_vmap
def tendon_vmap(unused_axis_size, is_batched, m, d): 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): def _mjx_contact(dx, worldid: int):
keys = [] keys = []
for i in range(dx._impl.ncon[0]): for i in range(dx._impl.nacon[0]):
if dx._impl.contact__worldid[i] != worldid: if dx._impl.contact__worldid[i] != worldid:
continue continue
mjx/mujoco/mjx/warp/types.py
+13 -23
View File
@@ -13,7 +13,6 @@
# limitations under the License. # limitations under the License.
# ============================================================================== # ==============================================================================
"""MJX Warp types. """MJX Warp types.
DO NOT EDIT. This file is auto-generated. DO NOT EDIT. This file is auto-generated.
""" """
import dataclasses import dataclasses
@@ -23,10 +22,8 @@ from jax import tree_util
from jax.interpreters import batching from jax.interpreters import batching
from mujoco.mjx._src import dataclasses as mjx_dataclasses from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np import numpy as np
PyTreeNode = mjx_dataclasses.PyTreeNode PyTreeNode = mjx_dataclasses.PyTreeNode
@dataclasses.dataclass(frozen=True) @dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class @tree_util.register_pytree_node_class
class TileSet: class TileSet:
@@ -38,7 +35,6 @@ class TileSet:
adr: address of each tile in the set adr: address of each tile in the set
size: size of all the tiles in this set size: size of all the tiles in this set
""" """
adr: np.ndarray adr: np.ndarray
size: int size: int
@@ -59,7 +55,6 @@ class BlockDim:
TODO(team): experimental and may be removed TODO(team): experimental and may be removed
""" """
actuator_velocity: int actuator_velocity: int
cholesky_factorize: int cholesky_factorize: int
cholesky_factorize_solve: int cholesky_factorize_solve: int
@@ -87,13 +82,10 @@ class BlockDim:
class StatisticWarp(PyTreeNode): class StatisticWarp(PyTreeNode):
"""Derived fields from Statistic.""" """Derived fields from Statistic."""
meaninertia: float meaninertia: float
class OptionWarp(PyTreeNode): class OptionWarp(PyTreeNode):
"""Derived fields from Option.""" """Derived fields from Option."""
broadphase: int broadphase: int
broadphase_filter: int broadphase_filter: int
ccd_iterations: int ccd_iterations: int
@@ -109,10 +101,8 @@ class OptionWarp(PyTreeNode):
sdf_initpoints: int sdf_initpoints: int
sdf_iterations: int sdf_iterations: int
class ModelWarp(PyTreeNode): class ModelWarp(PyTreeNode):
"""Derived fields from Model.""" """Derived fields from Model."""
M_colind: np.ndarray M_colind: np.ndarray
M_rowadr: np.ndarray M_rowadr: np.ndarray
M_rownnz: np.ndarray M_rownnz: np.ndarray
@@ -120,6 +110,7 @@ class ModelWarp(PyTreeNode):
actuator_moment_tiles_nv: Tuple[TileSet, ...] actuator_moment_tiles_nv: Tuple[TileSet, ...]
actuator_trntype_body_adr: np.ndarray actuator_trntype_body_adr: np.ndarray
block_dim: BlockDim block_dim: BlockDim
body_fluid_ellipsoid: np.ndarray
body_tree: Tuple[np.ndarray, ...] body_tree: Tuple[np.ndarray, ...]
collision_sensor_adr: np.ndarray collision_sensor_adr: np.ndarray
condim_max: int condim_max: int
@@ -218,10 +209,8 @@ class ModelWarp(PyTreeNode):
wrap_site_adr: np.ndarray wrap_site_adr: np.ndarray
wrap_site_pair_adr: np.ndarray wrap_site_pair_adr: np.ndarray
class DataWarp(PyTreeNode): class DataWarp(PyTreeNode):
"""Derived fields from Data.""" """Derived fields from Data."""
act_dot_rk: jax.Array act_dot_rk: jax.Array
act_t0: jax.Array act_t0: jax.Array
act_vel_integration: jax.Array act_vel_integration: jax.Array
@@ -317,9 +306,9 @@ class DataWarp(PyTreeNode):
multiccd_pdist: jax.Array multiccd_pdist: jax.Array
multiccd_pnormal: jax.Array multiccd_pnormal: jax.Array
multiccd_polygon: jax.Array multiccd_polygon: jax.Array
nacon: jax.Array
naconmax: int
ncollision: jax.Array ncollision: jax.Array
ncon: jax.Array
nconmax: int
ne: jax.Array ne: jax.Array
ne_connect: jax.Array ne_connect: jax.Array
ne_jnt: jax.Array ne_jnt: jax.Array
@@ -378,8 +367,6 @@ class DataWarp(PyTreeNode):
wrap_obj: jax.Array wrap_obj: jax.Array
wrap_xpos: jax.Array wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape) shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = { DATA_NON_VMAP = {
'collision_pair', 'collision_pair',
'collision_pairid', 'collision_pairid',
@@ -419,16 +406,15 @@ DATA_NON_VMAP = {
'multiccd_pdist', 'multiccd_pdist',
'multiccd_pnormal', 'multiccd_pnormal',
'multiccd_polygon', 'multiccd_polygon',
'nacon',
'naconmax',
'ncollision', 'ncollision',
'ncon',
'nconmax',
'njmax', 'njmax',
'nsolving', 'nsolving',
'nworld', 'nworld',
'ray_bodyexclude', 'ray_bodyexclude',
} }
def _to_elt(cont, _, d, axis): def _to_elt(cont, _, d, axis):
return DataWarp(**{ return DataWarp(**{
f.name: ( f.name: (
@@ -562,9 +548,9 @@ _NDIM = {
'multiccd_pdist': 2, 'multiccd_pdist': 2,
'multiccd_pnormal': 3, 'multiccd_pnormal': 3,
'multiccd_polygon': 3, 'multiccd_polygon': 3,
'nacon': 1,
'naconmax': 0,
'ncollision': 1, 'ncollision': 1,
'ncon': 1,
'nconmax': 0,
'ne': 1, 'ne': 1,
'ne_connect': 1, 'ne_connect': 1,
'ne_jnt': 1, 'ne_jnt': 1,
@@ -697,6 +683,7 @@ _NDIM = {
'body_contype': 1, 'body_contype': 1,
'body_dofadr': 1, 'body_dofadr': 1,
'body_dofnum': 1, 'body_dofnum': 1,
'body_fluid_ellipsoid': 1,
'body_geomadr': 1, 'body_geomadr': 1,
'body_geomnum': 1, 'body_geomnum': 1,
'body_gravcomp': 2, 'body_gravcomp': 2,
@@ -774,6 +761,7 @@ _NDIM = {
'geom_condim': 1, 'geom_condim': 1,
'geom_contype': 1, 'geom_contype': 1,
'geom_dataid': 1, 'geom_dataid': 1,
'geom_fluid': 2,
'geom_friction': 3, 'geom_friction': 3,
'geom_gap': 2, 'geom_gap': 2,
'geom_group': 1, 'geom_group': 1,
@@ -1155,9 +1143,9 @@ _BATCH_DIM = {
'multiccd_pdist': False, 'multiccd_pdist': False,
'multiccd_pnormal': False, 'multiccd_pnormal': False,
'multiccd_polygon': False, 'multiccd_polygon': False,
'nacon': False,
'naconmax': False,
'ncollision': False, 'ncollision': False,
'ncon': False,
'nconmax': False,
'ne': True, 'ne': True,
'ne_connect': True, 'ne_connect': True,
'ne_jnt': True, 'ne_jnt': True,
@@ -1290,6 +1278,7 @@ _BATCH_DIM = {
'body_contype': False, 'body_contype': False,
'body_dofadr': False, 'body_dofadr': False,
'body_dofnum': False, 'body_dofnum': False,
'body_fluid_ellipsoid': False,
'body_geomadr': False, 'body_geomadr': False,
'body_geomnum': False, 'body_geomnum': False,
'body_gravcomp': True, 'body_gravcomp': True,
@@ -1367,6 +1356,7 @@ _BATCH_DIM = {
'geom_condim': False, 'geom_condim': False,
'geom_contype': False, 'geom_contype': False,
'geom_dataid': False, 'geom_dataid': False,
'geom_fluid': False,
'geom_friction': True, 'geom_friction': True,
'geom_gap': True, 'geom_gap': True,
'geom_group': False, 'geom_group': False,