Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 929061622
Change-Id: I53892a819c3727b1f04db943839606ec1f477df1

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':

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

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]
+        acc += flex_bending[bendingadr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
-      force[i, x] = -(acc + flex_bending[adr + 16] * frc[i, x])
+      force[i, x] = -(acc + flex_bending[bendingadr + 16] * frc[i, x])
 
   for i in range(nvert):
     bodyid = flex_vertbodyid[v[i]]

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,9 +1442,6 @@ 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)
 
 
@@ -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()
+  conid_start, pairid = wp.tid()
-
-  dof1id = dof_tri_row[elementid]
-  dof2id = dof_tri_col[elementid]
 
   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
+    # One thread per (contact, support-pair): the support dofs are exactly the colind entries,
-    # the sparse positions of dof1id and dof2id. Skip if either is absent.
+    # 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]
+    npairs = rownnz * (rownnz + 1) // 2
-    pos1 = int(-1)
+    if pairid >= npairs:
-    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:
       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,8 +2550,16 @@ def _JTDAJ_sparse(
 
 def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
   # grad = Ma - qfrc_smooth - qfrc_constraint
+  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),
@@ -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
+      # Block-limit the launch: cap the grid near SM-filling width and stride over contacts, so
-      # loop unrolling
+      # 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)
+  _linesearch(m, d, ctx)
-
-  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],
-    )
 
   # 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,19 +3009,43 @@ def _solver_iteration(
       dim=d.nworld,
       outputs=[ctx.beta, ctx.beta_den],
     )
-    wp.launch(
+    wp.launch_tiled(
-      _solve_beta_accumulate,
+      _solve_beta_accumulate_tiled,
-      dim=(d.nworld, m.nv),
+      dim=d.nworld,
-      inputs=[ctx.grad, ctx.Mgrad, ctx.prev_grad, ctx.prev_Mgrad, ctx.done],
+      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_beta_finalize,
+      _solve_cg_finalize,
       dim=d.nworld,
-      inputs=[ctx.beta, ctx.beta_den, ctx.done],
+      inputs=[
-      outputs=[ctx.beta],
+        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,
     )
 
+  else:
     wp.launch(_solve_zero_search_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.search_dot])
 
     wp.launch(
@@ -3369,6 +3140,16 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
   init_context(m, d, ctx, grad=True)
 
   # search = -Mgrad
+  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,
+    )
+
+  else:
     wp.launch(
       _solve_init_search,
       dim=(d.nworld, m.nv),
@@ -3376,8 +3157,6 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
       outputs=[ctx.search, ctx.search_dot],
     )
 
-  step_size_cost = wp.empty((d.nworld, m.opt.ls_iterations if m.opt.ls_parallel else 0), dtype=float)
-
   nsolving = wp.full(shape=(1,), value=d.nworld, dtype=int)
   if m.opt.iterations != 0 and m.opt.graph_conditional:
     # Note: the iteration kernel (indicated by while_body) is repeatedly launched
@@ -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(
+    wp.capture_while(nsolving, while_body=_solver_iteration, m=m, d=d, ctx=ctx, nsolving=nsolving)
-      nsolving, while_body=_solver_iteration, m=m, d=d, ctx=ctx, step_size_cost=step_size_cost, 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])
   )

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]

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"

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,

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_cholesky_blocked: update gradient Cholesky blocked block
-    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
+      dimension (solver)
-    update_gradient_JTDAJ_dense: update gradient JTDAJ dense 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,