diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index f41bf332..1dec6057 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -182,8 +182,6 @@ class ModelIOTest(parameterized.TestCase): self.assertFalse(hasattr(mx, 'bvh_aabb')) elif impl == 'warp': - # Options specific to Warp are populated. - self.assertTrue(hasattr(mx.opt._impl, 'ls_parallel')) # Fields private to Warp backend impl are populated. self.assertTrue(hasattr(mx._impl, 'nxn_geom_pair')) elif impl == 'cpp': 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 92d0f7a1..c90451ab 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -257,9 +257,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: opt.tolerance = max(opt.tolerance, 1e-6) # warp only fields - ls_parallel_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "ls_parallel") - opt.ls_parallel = (ls_parallel_id > -1) and (mjm.numeric_data[mjm.numeric_adr[ls_parallel_id]] == 1) - opt.ls_parallel_min_step = 1.0e-6 # TODO(team): determine good default setting opt.broadphase = types.BroadphaseType.NXN opt.broadphase_filter = types.BroadphaseFilter.PLANE | types.BroadphaseFilter.SPHERE | types.BroadphaseFilter.OBB opt.graph_conditional = True @@ -301,6 +298,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1)) m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any() m.block_dim = types.BlockDim() + # Derive CG solver block_dim from nv: clamp(round_up_to_32(nv), 32, 256) + _nv_block = max(32, min(256, ((mjm.nv + 31) // 32) * 32)) + m.block_dim.update_gradient_grad = _nv_block + m.block_dim.solve_beta_accumulate = _nv_block + m.block_dim.solve_search_update_cg = _nv_block + m.block_dim.solve_init_search_cg = _nv_block if mjm.nv > 500: m.block_dim.linesearch_iterative = 512 m.is_sparse = is_sparse(mjm) @@ -308,6 +311,26 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.max_ten_J_rownnz = int(mjm.ten_J_rownnz.max()) if mjm.ntendon else 0 + # Upper bound on a contact's Jacobian support, to size the elliptic-cone JTCJ launch (one + # thread per (contact, support-pair)). A contact's row spans the dof chains of weld(b1) and + # weld(b2) (see _efc_contact_jac_sparse in constraint.py); take the largest union over all + # geom-carrying bodies -- a safe superset, since over-estimating only adds skipped threads. + # A body's dof chain is exactly the sparsity of its deepest dof's row in the (ancestor- + # structured) mass matrix, so reuse MuJoCo's precomputed M_colind rather than re-walking. + def _dof_chain(body): + if mjm.body_dofnum[body] == 0: + return frozenset() + dof = int(mjm.body_dofadr[body] + mjm.body_dofnum[body] - 1) + adr = int(mjm.M_rowadr[dof]) + return frozenset(int(mjm.M_colind[adr + k]) for k in range(int(mjm.M_rownnz[dof]))) + + chains = list({_dof_chain(int(mjm.body_weldid[b])) for b in mjm.geom_bodyid}) + max_rownnz = 0 + for i, chain_i in enumerate(chains): + for chain_j in chains[i:]: + max_rownnz = max(max_rownnz, len(chain_i | chain_j)) + m.jtcj_max_pairs = max(max_rownnz * (max_rownnz + 1) // 2, 1) + # body ids grouped by tree level (depth-based traversal) bodies, body_depth = {}, np.zeros(mjm.nbody, dtype=int) - 1 for i in range(mjm.nbody): @@ -2841,7 +2864,6 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any Overrides are of the format: opt.iterations = 1 - opt.ls_parallel = True opt.cone = pyramidal opt.disableflags = contact | spring """ @@ -2865,7 +2887,6 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any mjw_only_fields = { "opt.broadphase", "opt.broadphase_filter", - "opt.ls_parallel", "opt.graph_conditional", "opt.contact_sensor_maxmatch", } @@ -2881,6 +2902,11 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any overrides = overrides_dict for key, val in overrides.items(): + if key == "opt.ls_parallel": + raise ValueError("ls_parallel was removed in MuJoCo Warp 3.9.1.") + if key == "opt.ls_parallel_min_step": + raise ValueError("ls_parallel_min_step was removed in MuJoCo Warp 3.9.1.") + # skip overrides on MjModel for properties that are only on mjw.Model if key in mjw_only_fields and isinstance(model, mujoco.MjModel): continue 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 2d261abb..a1893ff6 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py @@ -630,6 +630,12 @@ def _flex_elasticity( f = i break + stiffness_adr_base = flex_stiffnessadr[f] + if stiffness_adr_base < 0: + return + if flex_stiffness[stiffness_adr_base] == 0.0: + return + local_elemid = elemid - flex_elemadr[f] dim = flex_dim[f] nvert = dim + 1 @@ -671,7 +677,7 @@ def _flex_elasticity( metric = wp.matrix(0.0, shape=(6, 6)) stiffness_size = nedge * (nedge + 1) / 2 - stiffness_adr = flex_stiffnessadr[f] + local_elemid * stiffness_size + stiffness_adr = stiffness_adr_base + local_elemid * stiffness_size id = int(0) for ed1 in range(nedge): for ed2 in range(ed1, nedge): @@ -721,6 +727,10 @@ def _flex_bending( f = i break + bendingadr = flex_bendingadr[f] + if bendingadr < 0: + return + if flex_dim[f] != 2: return @@ -734,10 +744,8 @@ def _flex_bending( flex_vertadr[f] + flex_edgeflap[edgeid][1], ) - adr = flex_bendingadr[f] - frc = wp.matrix(0.0, shape=(4, 3)) - if flex_bending[adr + 16]: + if flex_bending[bendingadr + 16]: v0 = flexvert_xpos_in[worldid, v[0]] v1 = flexvert_xpos_in[worldid, v[1]] v2 = flexvert_xpos_in[worldid, v[2]] @@ -752,8 +760,8 @@ def _flex_bending( for x in range(3): acc = float(0.0) for j in range(nvert): - acc += flex_bending[adr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x] - force[i, x] = -(acc + flex_bending[adr + 16] * frc[i, x]) + acc += flex_bending[bendingadr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x] + force[i, x] = -(acc + flex_bending[bendingadr + 16] * frc[i, x]) for i in range(nvert): bodyid = flex_vertbodyid[v[i]] 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 0329924a..a8459410 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py @@ -130,7 +130,6 @@ def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext: mv=wp.empty((nworld, nv), dtype=float), jv=wp.empty((nworld, njmax), dtype=float), quad=wp.empty((nworld, njmax), dtype=wp.vec3), - quad_gauss=wp.empty((nworld,), dtype=wp.vec3), alpha=wp.empty((nworld,), dtype=float), improvement=wp.empty((nworld,), dtype=float), prev_grad=wp.empty((nworld, nv), dtype=float), @@ -459,227 +458,6 @@ def _eval_constraint( return wp.vec3(-D * jaref, float(types.ConstraintState.QUADRATIC.value), 0.5 * D * jaref * jaref) -@wp.func -def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> float: - step = (wp.log(max_value) - wp.log(min_value)) / wp.max(1.0, float(num_values - 1)) - return wp.exp(wp.log(min_value) + float(i) * step) - - -@wp.kernel -def _linesearch_parallel_fused( - # Model: - opt_ls_iterations: int, - opt_impratio_invsqrt: wp.array[float], - opt_ls_parallel_min_step: float, - # Data in: - ne_in: wp.array[int], - nf_in: wp.array[int], - nefc_in: wp.array[int], - contact_friction_in: wp.array[types.vec5], - contact_efc_address_in: wp.array2d[int], - efc_type_in: wp.array2d[int], - efc_id_in: wp.array2d[int], - efc_D_in: wp.array2d[float], - efc_frictionloss_in: wp.array2d[float], - njmax_in: int, - nacon_in: wp.array[int], - # In: - ctx_Jaref_in: wp.array2d[float], - ctx_jv_in: wp.array2d[float], - ctx_quad_in: wp.array2d[wp.vec3], - ctx_quad_gauss_in: wp.array[wp.vec3], - ctx_done_in: wp.array[bool], - # Out: - cost_out: wp.array2d[float], -): - worldid, alphaid = wp.tid() - - if ctx_done_in[worldid]: - return - - alpha = _log_scale(opt_ls_parallel_min_step, 1.0, opt_ls_iterations, alphaid) - - quad_gauss = ctx_quad_gauss_in[worldid] - out = alpha * alpha * quad_gauss[2] + alpha * quad_gauss[1] - - ne = ne_in[worldid] - nf = nf_in[worldid] - - for efcid in range(min(njmax_in, nefc_in[worldid])): - # equality - if efcid < ne: - quad = ctx_quad_in[worldid, efcid] - out += alpha * alpha * quad[2] + alpha * quad[1] - # friction - elif efcid < ne + nf: - # search point, friction loss, bound (rf) - start = ctx_Jaref_in[worldid, efcid] - dir = ctx_jv_in[worldid, efcid] - x = start + alpha * dir - f = efc_frictionloss_in[worldid, efcid] - efc_D = efc_D_in[worldid, efcid] - rf = math.safe_div(f, efc_D) - - out += _eval_frictionloss_cost(x, f, rf, efc_D) - _eval_frictionloss_cost(start, f, rf, efc_D) - # limit and contact - elif efc_type_in[worldid, efcid] == types.ConstraintType.CONTACT_ELLIPTIC: - # extract contact info - conid = efc_id_in[worldid, efcid] - - if conid >= nacon_in[0]: - continue - - efcid0 = contact_efc_address_in[conid, 0] - if efcid != efcid0: - continue - - # unpack quad - efcid1 = contact_efc_address_in[conid, 1] - efcid2 = contact_efc_address_in[conid, 2] - - impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] - friction = contact_friction_in[conid] - quad = ctx_quad_in[worldid, efcid] - quad1 = ctx_quad_in[worldid, efcid1] - quad2 = ctx_quad_in[worldid, efcid2] - - mu = friction[0] * impratio_invsqrt - out += _eval_elliptic_cost(mu, quad, quad1, quad2, alpha) - out -= _eval_elliptic_cost(mu, quad, quad1, quad2, 0.0) - else: - # search point - start = ctx_Jaref_in[worldid, efcid] - x = start + alpha * ctx_jv_in[worldid, efcid] - cost0 = wp.where(start < 0.0, ctx_quad_in[worldid, efcid][0], 0.0) - - # active - if x < 0.0: - out += _eval_cost(ctx_quad_in[worldid, efcid], alpha) - cost0 - else: - out -= cost0 - - cost_out[worldid, alphaid] = out - - -@wp.kernel -def _linesearch_parallel_best_alpha( - # Model: - opt_ls_iterations: int, - opt_ls_parallel_min_step: float, - # In: - ctx_done_in: wp.array[bool], - cost_in: wp.array2d[float], - # Out: - ctx_alpha_out: wp.array[float], - ctx_improvement_out: wp.array[float], -): - worldid = wp.tid() - - if ctx_done_in[worldid]: - return - - bestid = int(0) - best_cost = float(0.0) - improved = bool(False) - for i in range(opt_ls_iterations): - cost = cost_in[worldid, i] - if cost < best_cost: - best_cost = cost - bestid = i - improved = True - - ctx_alpha_out[worldid] = wp.where(improved, _log_scale(opt_ls_parallel_min_step, 1.0, opt_ls_iterations, bestid), 0.0) - ctx_improvement_out[worldid] = wp.where(improved, -best_cost, 0.0) - - -def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d[float]): - """Parallel linesearch with setup and teardown kernels.""" - dofs_per_thread = 20 if m.nv > 50 else 50 - threads_per_efc = ceil(m.nv / dofs_per_thread) - - # quad_gauss = [0, search.T @ Ma - search.T @ qfrc_smooth, 0.5 * search.T @ mv] - if threads_per_efc > 1: - ctx.quad_gauss.zero_() - - wp.launch( - _linesearch_prepare_gauss(m.nv, dofs_per_thread), - dim=(d.nworld, threads_per_efc), - inputs=[d.qfrc_smooth, d.efc.Ma, ctx.search, ctx.mv, ctx.done], - outputs=[ctx.quad_gauss], - ) - - # quad = [0.5 * Jaref * Jaref * efc_D, jv * Jaref * efc_D, 0.5 * jv * jv * efc_D] - - wp.launch( - _linesearch_prepare_quad, - dim=(d.nworld, d.njmax), - inputs=[ - m.opt.impratio_invsqrt, - d.nefc, - d.contact.friction, - d.contact.dim, - d.contact.efc_address, - d.efc.type, - d.efc.id, - d.efc.D, - d.nacon, - ctx.Jaref, - ctx.jv, - ctx.done, - ], - outputs=[ctx.quad], - ) - - wp.launch( - _linesearch_parallel_fused, - dim=(d.nworld, m.opt.ls_iterations), - inputs=[ - m.opt.ls_iterations, - m.opt.impratio_invsqrt, - m.opt.ls_parallel_min_step, - d.ne, - d.nf, - d.nefc, - d.contact.friction, - d.contact.efc_address, - d.efc.type, - d.efc.id, - d.efc.D, - d.efc.frictionloss, - d.njmax, - d.nacon, - ctx.Jaref, - ctx.jv, - ctx.quad, - ctx.quad_gauss, - ctx.done, - ], - outputs=[cost], - ) - - wp.launch( - _linesearch_parallel_best_alpha, - dim=d.nworld, - inputs=[m.opt.ls_iterations, m.opt.ls_parallel_min_step, ctx.done, cost], - outputs=[ctx.alpha, ctx.improvement], - ) - - # Teardown: update qacc, Ma, Jaref - wp.launch( - _linesearch_qacc_ma, - dim=(d.nworld, m.nv), - inputs=[ctx.search, ctx.mv, ctx.alpha, ctx.done], - outputs=[d.qacc, d.efc.Ma], - ) - - wp.launch( - _linesearch_jaref, - dim=(d.nworld, d.njmax), - inputs=[d.nefc, ctx.jv, ctx.alpha, ctx.done], - outputs=[ctx.Jaref], - ) - - # kernel_analyzer: off @wp.func def _compute_efc_eval_pt_pyramidal( @@ -1628,205 +1406,21 @@ def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int): return kernel -@cache_kernel -def _linesearch_prepare_gauss(nv: int, dofs_per_thread: int): - @wp.kernel(module="unique", enable_backward=False) - def kernel( - # Data in: - qfrc_smooth_in: wp.array2d[float], - efc_Ma_in: wp.array2d[float], - # In: - ctx_search_in: wp.array2d[float], - ctx_mv_in: wp.array2d[float], - ctx_done_in: wp.array[bool], - # Out: - ctx_quad_gauss_out: wp.array[wp.vec3], - ): - worldid, dofstart = wp.tid() - - if ctx_done_in[worldid]: - return - - quad_gauss_1 = float(0.0) - quad_gauss_2 = float(0.0) - - if wp.static(dofs_per_thread >= nv): - for i in range(wp.static(nv)): - search = ctx_search_in[worldid, i] - quad_gauss_1 += search * (efc_Ma_in[worldid, i] - qfrc_smooth_in[worldid, i]) - quad_gauss_2 += 0.5 * search * ctx_mv_in[worldid, i] - - ctx_quad_gauss_out[worldid] = wp.vec3(0.0, quad_gauss_1, quad_gauss_2) - - else: - for i in range(wp.static(dofs_per_thread)): - ii = dofstart * wp.static(dofs_per_thread) + i - if ii < nv: - search = ctx_search_in[worldid, ii] - quad_gauss_1 += search * (efc_Ma_in[worldid, ii] - qfrc_smooth_in[worldid, ii]) - quad_gauss_2 += 0.5 * search * ctx_mv_in[worldid, ii] - - wp.atomic_add(ctx_quad_gauss_out, worldid, wp.vec3(0.0, quad_gauss_1, quad_gauss_2)) - - return kernel - - -@wp.kernel -def _linesearch_prepare_quad( - # Model: - opt_impratio_invsqrt: wp.array[float], - # Data in: - nefc_in: wp.array[int], - contact_friction_in: wp.array[types.vec5], - contact_dim_in: wp.array[int], - contact_efc_address_in: wp.array2d[int], - efc_type_in: wp.array2d[int], - efc_id_in: wp.array2d[int], - efc_D_in: wp.array2d[float], - nacon_in: wp.array[int], - # In: - ctx_Jaref_in: wp.array2d[float], - ctx_jv_in: wp.array2d[float], - ctx_done_in: wp.array[bool], - # Out: - ctx_quad_out: wp.array2d[wp.vec3], -): - worldid, efcid = wp.tid() - - if efcid >= nefc_in[worldid]: - return - - if ctx_done_in[worldid]: - return - - Jaref = ctx_Jaref_in[worldid, efcid] - jv = ctx_jv_in[worldid, efcid] - efc_D = efc_D_in[worldid, efcid] - - # init with scalar quadratic - quad = wp.vec3(0.5 * Jaref * Jaref * efc_D, jv * Jaref * efc_D, 0.5 * jv * jv * efc_D) - - # elliptic cone: extra processing - if efc_type_in[worldid, efcid] == types.ConstraintType.CONTACT_ELLIPTIC: - # extract contact info - conid = efc_id_in[worldid, efcid] - - if conid >= nacon_in[0]: - return - - efcid0 = contact_efc_address_in[conid, 0] - - if efcid != efcid0: - return - - dim = contact_dim_in[conid] - friction = contact_friction_in[conid] - mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] - - u0 = Jaref * mu - v0 = jv * mu - - uu = float(0.0) - uv = float(0.0) - vv = float(0.0) - for j in range(1, dim): - # complete vector quadratic (for bottom zone) - efcidj = contact_efc_address_in[conid, j] - if efcidj < 0: - return - jvj = ctx_jv_in[worldid, efcidj] - jarefj = ctx_Jaref_in[worldid, efcidj] - dj = efc_D_in[worldid, efcidj] - DJj = dj * jarefj - - quad += wp.vec3( - 0.5 * jarefj * DJj, - jvj * DJj, - 0.5 * jvj * dj * jvj, - ) - - # rescale to make primal cone circular - frictionj = friction[j - 1] - uj = jarefj * frictionj - vj = jvj * frictionj - - # accumulate sums of squares - uu += uj * uj - uv += uj * vj - vv += vj * vj - - quad1 = wp.vec3(u0, v0, uu) - efcid1 = contact_efc_address_in[conid, 1] - ctx_quad_out[worldid, efcid1] = quad1 - - mu2 = mu * mu - quad2 = wp.vec3(uv, vv, efc_D / (mu2 * (1.0 + mu2))) - efcid2 = contact_efc_address_in[conid, 2] - ctx_quad_out[worldid, efcid2] = quad2 - - ctx_quad_out[worldid, efcid] = quad - - -@wp.kernel -def _linesearch_qacc_ma( - # In: - ctx_search_in: wp.array2d[float], - ctx_mv_in: wp.array2d[float], - ctx_alpha_in: wp.array[float], - ctx_done_in: wp.array[bool], - # Data out: - qacc_out: wp.array2d[float], - efc_Ma_out: wp.array2d[float], -): - worldid, dofid = wp.tid() - - if ctx_done_in[worldid]: - return - - alpha = ctx_alpha_in[worldid] - qacc_out[worldid, dofid] += alpha * ctx_search_in[worldid, dofid] - efc_Ma_out[worldid, dofid] += alpha * ctx_mv_in[worldid, dofid] - - -@wp.kernel -def _linesearch_jaref( - # Data in: - nefc_in: wp.array[int], - # In: - ctx_jv_in: wp.array2d[float], - ctx_alpha_in: wp.array[float], - ctx_done_in: wp.array[bool], - # Out: - ctx_Jaref_out: wp.array2d[float], -): - worldid, efcid = wp.tid() - - if efcid >= nefc_in[worldid]: - return - - if ctx_done_in[worldid]: - return - - ctx_Jaref_out[worldid, efcid] += ctx_alpha_in[worldid] * ctx_jv_in[worldid, efcid] - - @event_scope -def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d[float]): +def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext): """Linesearch for constraint solver. Args: m: Model d: Data ctx: SolverContext - cost: Scratch array for storing costs per (world, alpha) - used for parallel mode """ # mv = M @ search (common to both parallel and iterative) support.mul_m(m, d, ctx.mv, ctx.search, skip=ctx.done) # Fuse jv computation in-kernel for small nv (iterative only, dense only) - # Parallel linesearch always requires jv pre-computed # Sparse mode requires pre-computed jv since in-kernel uses dense indexing - fuse_jv = m.nv <= 50 and not m.opt.ls_parallel and not m.is_sparse + fuse_jv = m.nv <= 50 and not m.is_sparse # jv = J @ search (when not fused into iterative kernel) if not fuse_jv: @@ -1848,10 +1442,7 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra outputs=[ctx.jv], ) - if m.opt.ls_parallel: - _linesearch_parallel(m, d, ctx, cost) - else: - _linesearch_iterative(m, d, ctx, fuse_jv) + _linesearch_iterative(m, d, ctx, fuse_jv) @wp.kernel @@ -1931,6 +1522,40 @@ def _solve_init_search( wp.atomic_add(ctx_search_dot_out, worldid, search * search) +@wp.kernel +def _solve_init_search_cg_tiled( + # Model: + nv: int, + # In: + ctx_grad_in: wp.array2d[float], + ctx_Mgrad_in: wp.array2d[float], + # Out: + ctx_search_out: wp.array2d[float], + ctx_search_dot_out: wp.array[float], + ctx_prev_grad_out: wp.array2d[float], + ctx_prev_Mgrad_out: wp.array2d[float], +): + worldid, tid = wp.tid() + + local_search_dot = float(0.0) + BLOCK_DIM = wp.block_dim() + + for dofid in range(tid, nv, BLOCK_DIM): + mgrad = ctx_Mgrad_in[worldid, dofid] + search = -1.0 * mgrad + ctx_search_out[worldid, dofid] = search + local_search_dot += search * search + + ctx_prev_grad_out[worldid, dofid] = ctx_grad_in[worldid, dofid] + ctx_prev_Mgrad_out[worldid, dofid] = mgrad + + search_dot_tile = wp.tile(local_search_dot, preserve_type=True) + search_dot_sum = wp.tile_reduce(wp.add, search_dot_tile) + + if tid == 0: + ctx_search_dot_out[worldid] = search_dot_sum[0] + + @cache_kernel def _update_constraint_efc(track_changes: bool): TRACK_CHANGES = track_changes @@ -2285,6 +1910,40 @@ def _update_gradient_grad( wp.atomic_add(ctx_grad_dot_out, worldid, grad * grad) +@wp.kernel +def _update_gradient_grad_tiled( + # Model: + nv: int, + # Data in: + qfrc_smooth_in: wp.array2d[float], + qfrc_constraint_in: wp.array2d[float], + efc_Ma_in: wp.array2d[float], + # In: + ctx_done_in: wp.array[bool], + # Out: + ctx_grad_out: wp.array2d[float], + ctx_grad_dot_out: wp.array[float], +): + worldid, tid = wp.tid() + + if ctx_done_in[worldid]: + return + + local_grad_dot = float(0.0) + BLOCK_DIM = wp.block_dim() + + for dofid in range(tid, nv, BLOCK_DIM): + grad = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - qfrc_constraint_in[worldid, dofid] + ctx_grad_out[worldid, dofid] = grad + local_grad_dot += grad * grad + + grad_dot_tile = wp.tile(local_grad_dot, preserve_type=True) + grad_dot_sum = wp.tile_reduce(wp.add, grad_dot_tile) + + if tid == 0: + ctx_grad_dot_out[worldid] = grad_dot_sum[0] + + @wp.kernel def _update_gradient_init_h_sparse( # Model: @@ -2400,8 +2059,6 @@ def _update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int): def _update_gradient_JTCJ_sparse( # Model: opt_impratio_invsqrt: wp.array[float], - dof_tri_row: wp.array[int], - dof_tri_col: wp.array[int], # Data in: contact_dist_in: wp.array[float], contact_includemargin_in: wp.array[float], @@ -2425,10 +2082,7 @@ def _update_gradient_JTCJ_sparse( # Out: ctx_h_out: wp.array3d[float], ): - conid_start, elementid = wp.tid() - - dof1id = dof_tri_row[elementid] - dof2id = dof_tri_col[elementid] + conid_start, pairid = wp.tid() for i in range(nblocks_perblock): conid = conid_start + i * dim_block @@ -2453,22 +2107,24 @@ def _update_gradient_JTCJ_sparse( if efc_state_in[worldid, efcid0] != types.ConstraintState.CONE: continue - # All dims share the same sparsity pattern. Scan colind once to find - # the sparse positions of dof1id and dof2id. Skip if either is absent. + # One thread per (contact, support-pair): the support dofs are exactly the colind entries, + # so decode pairid -> (pos1, pos2) with pos1 <= pos2 directly. No colind scan, and no + # membership skip (which the all-dof-pairs version wasted on ~99% absent dofs). rownnz = efc_J_rownnz_in[worldid, efcid0] - rowadr0 = efc_J_rowadr_in[worldid, efcid0] - pos1 = int(-1) - pos2 = int(-1) - for k in range(rownnz): - col = efc_J_colind_in[worldid, 0, rowadr0 + k] - if col == dof1id: - pos1 = k - if col == dof2id: - pos2 = k - if pos1 >= 0 and pos2 >= 0: - break - if pos1 < 0 or pos2 < 0: + npairs = rownnz * (rownnz + 1) // 2 + if pairid >= npairs: continue + rowadr0 = efc_J_rowadr_in[worldid, efcid0] + pos1 = int(0) + rem = pairid + while rem >= rownnz - pos1: + rem -= rownnz - pos1 + pos1 += 1 + pos2 = pos1 + rem + dofa = efc_J_colind_in[worldid, 0, rowadr0 + pos1] + dofb = efc_J_colind_in[worldid, 0, rowadr0 + pos2] + dof1id = wp.min(dofa, dofb) + dof2id = wp.max(dofa, dofb) fri = contact_friction_in[conid] mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] @@ -2555,7 +2211,8 @@ def _update_gradient_JTCJ_sparse( if dim1id != dim2id: h += hcone * efc_J12 * efc_J21 - ctx_h_out[worldid, dof1id, dof2id] += h + # multiple contacts can contribute to the same (dof1id, dof2id); atomic_add is exact + wp.atomic_add(ctx_h_out[worldid, dof1id], dof2id, h) @wp.kernel @@ -2893,14 +2550,22 @@ def _JTDAJ_sparse( def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): # grad = Ma - qfrc_smooth - qfrc_constraint - wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot]) - - wp.launch( - _update_gradient_grad, - dim=(d.nworld, m.nv), - inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done], - outputs=[ctx.grad, ctx.grad_dot], - ) + if m.opt.solver == types.SolverType.CG: + wp.launch_tiled( + _update_gradient_grad_tiled, + dim=d.nworld, + inputs=[m.nv, d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done], + outputs=[ctx.grad, ctx.grad_dot], + block_dim=m.block_dim.update_gradient_grad, + ) + else: + wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot]) + wp.launch( + _update_gradient_grad, + dim=(d.nworld, m.nv), + inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done], + outputs=[ctx.grad, ctx.grad_dot], + ) if m.opt.solver == types.SolverType.CG: smooth.solve_m(m, d, ctx.Mgrad, ctx.grad) @@ -2944,8 +2609,11 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): # of SMs on the GPU. We can now query the SM count: # https://github.com/NVIDIA/warp/commit/f3814e7e5459e5fd13032cf0fddb3daddd510f30 - # make dim_block and nblocks_perblock static for update_gradient_JTCJ to allow - # loop unrolling + # Block-limit the launch: cap the grid near SM-filling width and stride over contacts, so + # we don't over-launch naconmax (capacity) threads when active contacts are far fewer. The + # sparse kernel uses one thread per (contact, support-pair) (jtcj_max_pairs), the dense one + # per (contact, dof-pair) (dof_tri_row.size). + jtcj_second_dim = m.jtcj_max_pairs if m.is_sparse else m.dof_tri_row.size if wp.get_device().is_cuda: sm_count = wp.get_device().sm_count @@ -2953,7 +2621,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): # can be changed in the future to fine-tune the perf. The optimal factor will # depend on the kernel's occupancy, which determines how many blocks can # simultaneously run on the SM. TODO: This factor can be tuned further. - dim_block = ceil((sm_count * 6 * 256) / m.dof_tri_row.size) + dim_block = ceil((sm_count * 6 * 256) / jtcj_second_dim) else: # fall back for CPU dim_block = d.naconmax @@ -2963,11 +2631,9 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext): if m.is_sparse: wp.launch( _update_gradient_JTCJ_sparse, - dim=(dim_block, m.dof_tri_row.size), + dim=(dim_block, m.jtcj_max_pairs), inputs=[ m.opt.impratio_invsqrt, - m.dof_tri_row, - m.dof_tri_col, d.contact.dist, d.contact.includemargin, d.contact.friction, @@ -3071,25 +2737,6 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte _cholesky_factorize_solve(m, d, ctx, skip_unchanged=True) -@wp.kernel -def _solve_prev_grad_Mgrad( - # In: - ctx_grad_in: wp.array2d[float], - ctx_Mgrad_in: wp.array2d[float], - ctx_done_in: wp.array[bool], - # Out: - ctx_prev_grad_out: wp.array2d[float], - ctx_prev_Mgrad_out: wp.array2d[float], -): - worldid, dofid = wp.tid() - - if ctx_done_in[worldid]: - return - - ctx_prev_grad_out[worldid, dofid] = ctx_grad_in[worldid, dofid] - ctx_prev_Mgrad_out[worldid, dofid] = ctx_Mgrad_in[worldid, dofid] - - @wp.kernel def _solve_beta_zero( # Out: @@ -3101,6 +2748,47 @@ def _solve_beta_zero( ctx_beta_den_out[worldid] = 0.0 +@wp.kernel +def _solve_beta_accumulate_tiled( + # Model: + nv: int, + # In: + ctx_grad_in: wp.array2d[float], + ctx_Mgrad_in: wp.array2d[float], + ctx_prev_grad_in: wp.array2d[float], + ctx_prev_Mgrad_in: wp.array2d[float], + ctx_done_in: wp.array[bool], + # Out: + ctx_beta_num_out: wp.array[float], + ctx_beta_den_out: wp.array[float], +): + worldid, tid = wp.tid() + + if ctx_done_in[worldid]: + return + + local_num = float(0.0) + local_den = float(0.0) + BLOCK_DIM = wp.block_dim() + + for dofid in range(tid, nv, BLOCK_DIM): + prev_Mgrad = ctx_prev_Mgrad_in[worldid, dofid] + num = ctx_grad_in[worldid, dofid] * (ctx_Mgrad_in[worldid, dofid] - prev_Mgrad) + den = ctx_prev_grad_in[worldid, dofid] * prev_Mgrad + local_num += num + local_den += den + + num_tile = wp.tile(local_num, preserve_type=True) + num_sum = wp.tile_reduce(wp.add, num_tile) + + den_tile = wp.tile(local_den, preserve_type=True) + den_sum = wp.tile_reduce(wp.add, den_tile) + + if tid == 0: + ctx_beta_num_out[worldid] = num_sum[0] + ctx_beta_den_out[worldid] = den_sum[0] + + @wp.kernel def _solve_beta_accumulate( # In: @@ -3125,23 +2813,6 @@ def _solve_beta_accumulate( wp.atomic_add(ctx_beta_den_out, worldid, den) -@wp.kernel -def _solve_beta_finalize( - # In: - ctx_beta_num_in: wp.array[float], - ctx_beta_den_in: wp.array[float], - ctx_done_in: wp.array[bool], - # Out: - ctx_beta_out: wp.array[float], -): - worldid = wp.tid() - - if ctx_done_in[worldid]: - return - - ctx_beta_out[worldid] = wp.max(0.0, ctx_beta_num_in[worldid] / wp.max(types.MJ_MINVAL, ctx_beta_den_in[worldid])) - - @wp.kernel def _solve_zero_search_dot( # In: @@ -3184,6 +2855,91 @@ def _solve_search_update( wp.atomic_add(ctx_search_dot_out, worldid, search * search) +@wp.kernel +def _solve_search_update_cg_tiled( + # Model: + nv: int, + # In: + ctx_grad_in: wp.array2d[float], + ctx_Mgrad_in: wp.array2d[float], + ctx_search_in: wp.array2d[float], + ctx_beta_in: wp.array[float], + ctx_done_in: wp.array[bool], + # Out: + ctx_search_out: wp.array2d[float], + ctx_search_dot_out: wp.array[float], + ctx_prev_grad_out: wp.array2d[float], + ctx_prev_Mgrad_out: wp.array2d[float], +): + worldid, tid = wp.tid() + + if ctx_done_in[worldid]: + return + + local_search_dot = float(0.0) + BLOCK_DIM = wp.block_dim() + beta = ctx_beta_in[worldid] + + for dofid in range(tid, nv, BLOCK_DIM): + mgrad = ctx_Mgrad_in[worldid, dofid] + search = -1.0 * mgrad + beta * ctx_search_in[worldid, dofid] + + ctx_search_out[worldid, dofid] = search + local_search_dot += search * search + + ctx_prev_grad_out[worldid, dofid] = ctx_grad_in[worldid, dofid] + ctx_prev_Mgrad_out[worldid, dofid] = mgrad + + search_dot_tile = wp.tile(local_search_dot, preserve_type=True) + search_dot_sum = wp.tile_reduce(wp.add, search_dot_tile) + + if tid == 0: + ctx_search_dot_out[worldid] = search_dot_sum[0] + + +@wp.kernel +def _solve_cg_finalize( + # Model: + nv: int, + opt_tolerance: wp.array[float], + opt_iterations: int, + stat_meaninertia: wp.array[float], + # In: + ctx_beta_num_in: wp.array[float], + ctx_beta_den_in: wp.array[float], + ctx_improvement_in: wp.array[float], + ctx_done_in: wp.array[bool], + ctx_grad_dot_in: wp.array[float], + # Data out: + solver_niter_out: wp.array[int], + # Out: + ctx_beta_out: wp.array[float], + nsolving_out: wp.array[int], + ctx_done_out: wp.array[bool], +): + worldid = wp.tid() + + if ctx_done_in[worldid]: + return + + # 1. solve_beta_finalize + ctx_beta_out[worldid] = wp.max(0.0, ctx_beta_num_in[worldid] / wp.max(types.MJ_MINVAL, ctx_beta_den_in[worldid])) + + # 2. solve_done + solver_niter_out[worldid] += 1 + tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]] + meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]] + + grad_dot = ctx_grad_dot_in[worldid] + + improvement = _rescale(nv, meaninertia, ctx_improvement_in[worldid]) + gradient = _rescale(nv, meaninertia, wp.sqrt(grad_dot)) + done = (improvement < tolerance) or (gradient < tolerance) + if done or solver_niter_out[worldid] == opt_iterations: + ctx_done_out[worldid] = True + wp.atomic_add(nsolving_out, 0, -1) + + @wp.kernel def _solve_done( # Model: @@ -3225,18 +2981,9 @@ def _solver_iteration( m: types.Model, d: types.Data, ctx: SolverContext, - step_size_cost: wp.array2d[float], nsolving: wp.array[int], ): - _linesearch(m, d, ctx, step_size_cost) - - if m.opt.solver == types.SolverType.CG: - wp.launch( - _solve_prev_grad_Mgrad, - dim=(d.nworld, m.nv), - inputs=[ctx.grad, ctx.Mgrad, ctx.done], - outputs=[ctx.prev_grad, ctx.prev_Mgrad], - ) + _linesearch(m, d, ctx) # Incremental H is only valid for non-elliptic cones. The elliptic cone # path in _update_constraint_efc has early returns that skip state change @@ -3262,42 +3009,66 @@ def _solver_iteration( dim=d.nworld, outputs=[ctx.beta, ctx.beta_den], ) - wp.launch( - _solve_beta_accumulate, - dim=(d.nworld, m.nv), - inputs=[ctx.grad, ctx.Mgrad, ctx.prev_grad, ctx.prev_Mgrad, ctx.done], - outputs=[ctx.beta, ctx.beta_den], - ) - wp.launch( - _solve_beta_finalize, + wp.launch_tiled( + _solve_beta_accumulate_tiled, dim=d.nworld, - inputs=[ctx.beta, ctx.beta_den, ctx.done], - outputs=[ctx.beta], + inputs=[m.nv, ctx.grad, ctx.Mgrad, ctx.prev_grad, ctx.prev_Mgrad, ctx.done], + outputs=[ctx.beta, ctx.beta_den], + block_dim=m.block_dim.solve_beta_accumulate, + ) + wp.launch( + _solve_cg_finalize, + dim=d.nworld, + inputs=[ + m.nv, + m.opt.tolerance, + m.opt.iterations, + m.stat.meaninertia, + ctx.beta, + ctx.beta_den, + ctx.improvement, + ctx.done, + ctx.grad_dot, + ], + outputs=[ + d.solver_niter, + ctx.beta, + nsolving, + ctx.done, + ], + ) + wp.launch_tiled( + _solve_search_update_cg_tiled, + dim=d.nworld, + inputs=[m.nv, ctx.grad, ctx.Mgrad, ctx.search, ctx.beta, ctx.done], + outputs=[ctx.search, ctx.search_dot, ctx.prev_grad, ctx.prev_Mgrad], + block_dim=m.block_dim.solve_search_update_cg, ) - wp.launch(_solve_zero_search_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.search_dot]) + else: + wp.launch(_solve_zero_search_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.search_dot]) - wp.launch( - _solve_search_update, - dim=(d.nworld, m.nv), - inputs=[m.opt.solver, ctx.Mgrad, ctx.search, ctx.beta, ctx.done], - outputs=[ctx.search, ctx.search_dot], - ) + wp.launch( + _solve_search_update, + dim=(d.nworld, m.nv), + inputs=[m.opt.solver, ctx.Mgrad, ctx.search, ctx.beta, ctx.done], + outputs=[ctx.search, ctx.search_dot], + ) - wp.launch( - _solve_done, - dim=d.nworld, - inputs=[ - m.nv, - m.opt.tolerance, - m.opt.iterations, - m.stat.meaninertia, - ctx.grad_dot, - ctx.improvement, - ctx.done, - ], - outputs=[d.solver_niter, nsolving, ctx.done], - ) + wp.launch( + _solve_done, + dim=d.nworld, + inputs=[ + m.nv, + m.opt.tolerance, + m.opt.iterations, + m.stat.meaninertia, + ctx.grad_dot, + ctx.improvement, + ctx.done, + ], + outputs=[d.solver_niter, nsolving, ctx.done], + ) def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, grad: bool = True): @@ -3369,14 +3140,22 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext): init_context(m, d, ctx, grad=True) # search = -Mgrad - wp.launch( - _solve_init_search, - dim=(d.nworld, m.nv), - inputs=[ctx.Mgrad], - outputs=[ctx.search, ctx.search_dot], - ) + if m.opt.solver == types.SolverType.CG: + wp.launch_tiled( + _solve_init_search_cg_tiled, + dim=d.nworld, + inputs=[m.nv, ctx.grad, ctx.Mgrad], + outputs=[ctx.search, ctx.search_dot, ctx.prev_grad, ctx.prev_Mgrad], + block_dim=m.block_dim.solve_init_search_cg, + ) - step_size_cost = wp.empty((d.nworld, m.opt.ls_iterations if m.opt.ls_parallel else 0), dtype=float) + else: + wp.launch( + _solve_init_search, + dim=(d.nworld, m.nv), + inputs=[ctx.Mgrad], + outputs=[ctx.search, ctx.search_dot], + ) nsolving = wp.full(shape=(1,), value=d.nworld, dtype=int) if m.opt.iterations != 0 and m.opt.graph_conditional: @@ -3387,15 +3166,13 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext): # When the number of iterations reaches m.opt.iterations, solver_niter # becomes zero and all worlds are marked as converged to avoid an infinite loop. # note: we only launch the iteration kernel if everything is not done - wp.capture_while( - nsolving, while_body=_solver_iteration, m=m, d=d, ctx=ctx, step_size_cost=step_size_cost, nsolving=nsolving - ) + wp.capture_while(nsolving, while_body=_solver_iteration, m=m, d=d, ctx=ctx, nsolving=nsolving) else: # This branch is mostly for when JAX is used as it is currently not compatible # with CUDA graph conditional. # It should be removed when JAX becomes compatible. for _ in range(m.opt.iterations): - _solver_iteration(m, d, ctx, step_size_cost, nsolving) + _solver_iteration(m, d, ctx, nsolving) # TODO(team): Consolidate monolithic and island solver code where possible @@ -4510,6 +4287,10 @@ def _solve_beta_island_finalize( if islandid >= nisland_in[worldid]: return + if island_done_in[worldid, islandid]: + island_beta_out[worldid, islandid] = 0.0 + return + island_beta_out[worldid, islandid] = wp.max( 0.0, island_beta_num_in[worldid, islandid] / wp.max(types.MJ_MINVAL, island_beta_den_in[worldid, islandid]) ) 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 2fe4607d..941e1507 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -66,6 +66,10 @@ class BlockDim: update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver) update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver) linesearch_iterative: linesearch iterative block dimension (solver) + update_gradient_grad: update gradient grad block dimension (solver) + solve_beta_accumulate: solve beta accumulate block dimension (solver) + solve_search_update_cg: solve search update CG block dimension (solver) + solve_init_search_cg: solve init search CG block dimension (solver) contact_jac_tiled: contact Jacobian tiled block dimension (solver) qderiv_actuator_dense: qderiv actuator dense block dimension (derivative) render: render block dimension (render) @@ -92,6 +96,10 @@ class BlockDim: update_gradient_JTDAJ_sparse: int = 64 update_gradient_JTDAJ_dense: int = 128 linesearch_iterative: int = 32 + update_gradient_grad: int = 256 + solve_beta_accumulate: int = 256 + solve_search_update_cg: int = 256 + solve_init_search_cg: int = 256 contact_jac_tiled: int = 32 # derivative qderiv_actuator_dense: int = 32 @@ -800,8 +808,6 @@ class Option: warp only fields: impratio_invsqrt: ratio of friction-to-normal contact impedance (stored as inverse square root) - ls_parallel: evaluate engine solver step sizes in parallel - ls_parallel_min_step: minimum step size for solver linesearch broadphase: broadphase type (BroadphaseType) broadphase_filter: broadphase filter bitflag (BroadphaseFilter) graph_conditional: flag to use cuda graph conditional @@ -834,14 +840,29 @@ class Option: sdf_iterations: int # warp only fields: impratio_invsqrt: array("*", float) - ls_parallel: bool - ls_parallel_min_step: float broadphase: BroadphaseType broadphase_filter: BroadphaseFilter graph_conditional: bool run_collision_detection: bool contact_sensor_maxmatch: int + # TODO(team): remove in future version + @property + def ls_parallel(self) -> bool: + raise AttributeError("ls_parallel was removed in MuJoCo Warp 3.9.1.") + + @ls_parallel.setter + def ls_parallel(self, value: bool): + raise AttributeError("ls_parallel was removed in MuJoCo Warp 3.9.1.") + + @property + def ls_parallel_min_step(self) -> float: + raise AttributeError("ls_parallel_min_step was removed in MuJoCo Warp 3.9.1.") + + @ls_parallel_min_step.setter + def ls_parallel_min_step(self, value: float): + raise AttributeError("ls_parallel_min_step was removed in MuJoCo Warp 3.9.1.") + @dataclasses.dataclass class Statistic: @@ -1290,6 +1311,7 @@ class Model: tendon_geom_adr: geom tendon address tendon_limited_adr: addresses for limited tendons max_ten_J_rownnz: maximum number of non-zeros in a tendon row + jtcj_max_pairs: bound on a contact's support-pair count, sizes the elliptic-cone JTCJ launch ten_wrapadr_site: wrap object starting address for sites ten_wrapnum_site: number of site wrap objects per tendon wrap_jnt_adr: addresses for joint tendon wrap object @@ -1719,6 +1741,7 @@ class Model: tendon_geom_adr: wp.array[int] tendon_limited_adr: wp.array[int] max_ten_J_rownnz: int + jtcj_max_pairs: int ten_wrapadr_site: wp.array[int] ten_wrapnum_site: wp.array[int] wrap_jnt_adr: wp.array[int] @@ -2205,7 +2228,6 @@ class SolverContext: mv: wp.array2d[float] jv: wp.array2d[float] quad: wp.array2d[wp.vec3] - quad_gauss: wp.array[wp.vec3] alpha: wp.array[float] improvement: wp.array[float] prev_grad: wp.array2d[float] diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py index bac30dd3..5193c9a4 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py @@ -169,7 +169,7 @@ def _main(argv: Sequence[str]) -> None: solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name integrator = mjw.IntegratorType(m.opt.integrator).name iterations, ls_iterations = m.opt.iterations, m.opt.ls_iterations - ls_str = f"{'parallel' if m.opt.ls_parallel else 'iterative'} linesearch iterations: {ls_iterations}" + ls_str = f"linesearch iterations: {ls_iterations}" print( f" nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.is_sparse}\n" f" broadphase: {broadphase} broadphase_filter: {filter}\n" diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index a94949c8..c3238c02 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -14,17 +14,14 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" - import dataclasses import functools - import jax -import warp as wp - from mujoco.mjx._src import types +from mujoco.mjx.warp import ffi import mujoco.mjx.third_party.mujoco_warp as mjwarp from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types -from mujoco.mjx.warp import ffi +import warp as wp _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} @@ -241,6 +238,7 @@ def _forward_shim( jnt_stiffness: wp.array2d[float], jnt_stiffnesspoly: wp.array2d[wp.vec2], jnt_type: wp.array[int], + jtcj_max_pairs: int, light_bodyid: wp.array[int], light_dir: wp.array2d[wp.vec3], light_dir0: wp.array2d[wp.vec3], @@ -415,8 +413,6 @@ def _forward_shim( opt__integrator: int, opt__iterations: int, opt__ls_iterations: int, - opt__ls_parallel: bool, - opt__ls_parallel_min_step: float, opt__ls_tolerance: wp.array[float], opt__magnetic: wp.array[wp.vec3], opt__run_collision_detection: bool, @@ -780,6 +776,7 @@ def _forward_shim( _m.jnt_stiffness = jnt_stiffness _m.jnt_stiffnesspoly = jnt_stiffnesspoly _m.jnt_type = jnt_type + _m.jtcj_max_pairs = jtcj_max_pairs _m.light_bodyid = light_bodyid _m.light_dir = light_dir _m.light_dir0 = light_dir0 @@ -865,8 +862,6 @@ def _forward_shim( _m.opt.integrator = opt__integrator _m.opt.iterations = opt__iterations _m.opt.ls_iterations = opt__ls_iterations - _m.opt.ls_parallel = opt__ls_parallel - _m.opt.ls_parallel_min_step = opt__ls_parallel_min_step _m.opt.ls_tolerance = opt__ls_tolerance _m.opt.magnetic = opt__magnetic _m.opt.run_collision_detection = opt__run_collision_detection @@ -1796,6 +1791,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.jnt_stiffness, m.jnt_stiffnesspoly, m.jnt_type, + m._impl.jtcj_max_pairs, m._impl.light_bodyid, m.light_dir, m.light_dir0, @@ -1970,8 +1966,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): m.opt.integrator, m.opt.iterations, m.opt.ls_iterations, - m.opt._impl.ls_parallel, - m.opt._impl.ls_parallel_min_step, m.opt.ls_tolerance, m.opt.magnetic, m.opt._impl.run_collision_detection, @@ -2498,6 +2492,7 @@ def _step_shim( jnt_stiffness: wp.array2d[float], jnt_stiffnesspoly: wp.array2d[wp.vec2], jnt_type: wp.array[int], + jtcj_max_pairs: int, light_bodyid: wp.array[int], light_dir: wp.array2d[wp.vec3], light_dir0: wp.array2d[wp.vec3], @@ -2677,8 +2672,6 @@ def _step_shim( opt__integrator: int, opt__iterations: int, opt__ls_iterations: int, - opt__ls_parallel: bool, - opt__ls_parallel_min_step: float, opt__ls_tolerance: wp.array[float], opt__magnetic: wp.array[wp.vec3], opt__run_collision_detection: bool, @@ -3048,6 +3041,7 @@ def _step_shim( _m.jnt_stiffness = jnt_stiffness _m.jnt_stiffnesspoly = jnt_stiffnesspoly _m.jnt_type = jnt_type + _m.jtcj_max_pairs = jtcj_max_pairs _m.light_bodyid = light_bodyid _m.light_dir = light_dir _m.light_dir0 = light_dir0 @@ -3136,8 +3130,6 @@ def _step_shim( _m.opt.integrator = opt__integrator _m.opt.iterations = opt__iterations _m.opt.ls_iterations = opt__ls_iterations - _m.opt.ls_parallel = opt__ls_parallel - _m.opt.ls_parallel_min_step = opt__ls_parallel_min_step _m.opt.ls_tolerance = opt__ls_tolerance _m.opt.magnetic = opt__magnetic _m.opt.run_collision_detection = opt__run_collision_detection @@ -4089,6 +4081,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.jnt_stiffness, m.jnt_stiffnesspoly, m.jnt_type, + m._impl.jtcj_max_pairs, m._impl.light_bodyid, m.light_dir, m.light_dir0, @@ -4268,8 +4261,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): m.opt.integrator, m.opt.iterations, m.opt.ls_iterations, - m.opt._impl.ls_parallel, - m.opt._impl.ls_parallel_min_step, m.opt.ls_tolerance, m.opt.magnetic, m.opt._impl.run_collision_detection, diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index dc29da96..6e953948 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -15,17 +15,14 @@ """MJX Warp types. DO NOT EDIT. This file is auto-generated. """ - import dataclasses import typing from typing import Tuple - import jax from jax import tree_util from jax.interpreters import batching -import numpy as np - from mujoco.mjx._src import dataclasses as mjx_dataclasses +import numpy as np if typing.TYPE_CHECKING: GraphMode = int @@ -37,7 +34,6 @@ if typing.TYPE_CHECKING: else: try: from warp._src.jax_experimental.ffi import GraphMode - from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types Callback = mjwp_types.Callback @@ -46,7 +42,6 @@ else: Callback = None PyTreeNode = mjx_dataclasses.PyTreeNode - @dataclasses.dataclass(frozen=True) @tree_util.register_pytree_node_class class TileSet: @@ -58,12 +53,9 @@ class TileSet: adr: address of each tile in the set size: size of all the tiles in this set """ - adr: np.ndarray size: int - # Manually kept in this generated shim until TileSet method generation is - # needed more broadly. Keep this in sync with mujoco_warp._src.types.TileSet. def __eq__(self, other) -> bool: if self.__class__ is not other.__class__: return NotImplemented @@ -101,18 +93,25 @@ class BlockDim: energy_vel_kinetic: energy velocity kinetic block dimension (sensor) cholesky_factorize: Cholesky factorize block dimension (smooth) cholesky_solve: Cholesky solve block dimension (smooth) - cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth) + cholesky_factorize_solve: Cholesky factorize and solve block dimension + (smooth) solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth) update_gradient_cholesky: update gradient Cholesky block dimension (solver) - update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver) - update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver) - update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver) + update_gradient_cholesky_blocked: update gradient Cholesky blocked block + dimension (solver) + update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension + (solver) + update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension + (solver) linesearch_iterative: linesearch iterative block dimension (solver) + update_gradient_grad: update gradient grad block dimension (solver) + solve_beta_accumulate: solve beta accumulate block dimension (solver) + solve_search_update_cg: solve search update CG block dimension (solver) + solve_init_search_cg: solve init search CG block dimension (solver) contact_jac_tiled: contact Jacobian tiled block dimension (solver) qderiv_actuator_dense: qderiv actuator dense block dimension (derivative) render: render block dimension (render) """ - actuator_velocity: int cholesky_factorize: int cholesky_factorize_solve: int @@ -127,10 +126,14 @@ class BlockDim: render: int segmented_sort: int solve_LD_sparse_fused: int + solve_beta_accumulate: int + solve_init_search_cg: int + solve_search_update_cg: int update_gradient_JTDAJ_dense: int update_gradient_JTDAJ_sparse: int update_gradient_cholesky: int update_gradient_cholesky_blocked: int + update_gradient_grad: int def tree_flatten(self): children = list((getattr(self, k) for k in self.__dataclass_fields__)) @@ -144,13 +147,10 @@ class BlockDim: class StatisticWarp(PyTreeNode): """Derived fields from Statistic.""" - meaninertia: jax.Array - class OptionWarp(PyTreeNode): """Derived fields from Option.""" - broadphase: int broadphase_filter: int ccd_iterations: int @@ -159,17 +159,13 @@ class OptionWarp(PyTreeNode): graph_conditional: bool graph_mode: GraphMode impratio_invsqrt: jax.Array - ls_parallel: bool - ls_parallel_min_step: float run_collision_detection: bool sdf_initpoints: int sdf_iterations: int sleep_tolerance: jax.Array - class ModelWarp(PyTreeNode): """Derived fields from Model.""" - D_colind: np.ndarray D_diag: np.ndarray D_rowadr: np.ndarray @@ -254,6 +250,7 @@ class ModelWarp(PyTreeNode): is_sparse: bool jnt_limited_ball_adr: np.ndarray jnt_limited_slide_hinge_adr: np.ndarray + jtcj_max_pairs: int light_bodyid: np.ndarray light_targetbodyid: np.ndarray mapD2M: np.ndarray @@ -351,10 +348,8 @@ class ModelWarp(PyTreeNode): wrap_site_adr: np.ndarray wrap_site_pair_adr: np.ndarray - class DataWarp(PyTreeNode): """Derived fields from Data.""" - M: jax.Array actuator_moment: jax.Array actuator_velocity: jax.Array @@ -473,8 +468,6 @@ class DataWarp(PyTreeNode): wrap_obj: jax.Array wrap_xpos: jax.Array shape = property(lambda self: self.cacc.shape) - - DATA_NON_VMAP = { 'contact__dim', 'contact__dist', @@ -503,7 +496,6 @@ DATA_NON_VMAP = { 'nworld', } - def _to_elt(cont, _, d, axis): return DataWarp(**{ f.name: ( @@ -749,10 +741,14 @@ _NDIM = { 'block_dim__render': 0, 'block_dim__segmented_sort': 0, 'block_dim__solve_LD_sparse_fused': 0, + 'block_dim__solve_beta_accumulate': 0, + 'block_dim__solve_init_search_cg': 0, + 'block_dim__solve_search_update_cg': 0, 'block_dim__update_gradient_JTDAJ_dense': 0, 'block_dim__update_gradient_JTDAJ_sparse': 0, 'block_dim__update_gradient_cholesky': 0, 'block_dim__update_gradient_cholesky_blocked': 0, + 'block_dim__update_gradient_grad': 0, 'body_branch_start': 1, 'body_branches': 1, 'body_conaffinity': 1, @@ -914,6 +910,7 @@ _NDIM = { 'jnt_stiffness': 2, 'jnt_stiffnesspoly': 3, 'jnt_type': 1, + 'jtcj_max_pairs': 0, 'light_active': 2, 'light_ambient': 3, 'light_attenuation': 3, @@ -1043,8 +1040,6 @@ _NDIM = { 'opt__integrator': 0, 'opt__iterations': 0, 'opt__ls_iterations': 0, - 'opt__ls_parallel': 0, - 'opt__ls_parallel_min_step': 0, 'opt__ls_tolerance': 1, 'opt__magnetic': 2, 'opt__run_collision_detection': 0, @@ -1170,8 +1165,6 @@ _NDIM = { 'integrator': 0, 'iterations': 0, 'ls_iterations': 0, - 'ls_parallel': 0, - 'ls_parallel_min_step': 0, 'ls_tolerance': 1, 'magnetic': 2, 'run_collision_detection': 0, @@ -1407,10 +1400,14 @@ _BATCH_DIM = { 'block_dim__render': False, 'block_dim__segmented_sort': False, 'block_dim__solve_LD_sparse_fused': False, + 'block_dim__solve_beta_accumulate': False, + 'block_dim__solve_init_search_cg': False, + 'block_dim__solve_search_update_cg': False, 'block_dim__update_gradient_JTDAJ_dense': False, 'block_dim__update_gradient_JTDAJ_sparse': False, 'block_dim__update_gradient_cholesky': False, 'block_dim__update_gradient_cholesky_blocked': False, + 'block_dim__update_gradient_grad': False, 'body_branch_start': False, 'body_branches': False, 'body_conaffinity': False, @@ -1572,6 +1569,7 @@ _BATCH_DIM = { 'jnt_stiffness': True, 'jnt_stiffnesspoly': True, 'jnt_type': False, + 'jtcj_max_pairs': False, 'light_active': True, 'light_ambient': True, 'light_attenuation': True, @@ -1701,8 +1699,6 @@ _BATCH_DIM = { 'opt__integrator': False, 'opt__iterations': False, 'opt__ls_iterations': False, - 'opt__ls_parallel': False, - 'opt__ls_parallel_min_step': False, 'opt__ls_tolerance': True, 'opt__magnetic': True, 'opt__run_collision_detection': False, @@ -1828,8 +1824,6 @@ _BATCH_DIM = { 'integrator': False, 'iterations': False, 'ls_iterations': False, - 'ls_parallel': False, - 'ls_parallel_min_step': False, 'ls_tolerance': True, 'magnetic': True, 'run_collision_detection': False,