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