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

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 929061622
Change-Id: I53892a819c3727b1f04db943839606ec1f477df1
This commit is contained in:
Taylor Howell
2026-06-09 02:27:41 -07:00
committed by Copybara-Service
parent 58a8c6ee56
commit 763e713d1b
8 changed files with 425 additions and 605 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io_test.py
-2
View File
@@ -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
Vendored
+31 -5
View File
@@ -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
Vendored
+14 -6
View File
@@ -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]]
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+279 -498
View File
@@ -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()
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,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
# 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,19 +3009,43 @@ 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],
wp.launch_tiled(
_solve_beta_accumulate_tiled,
dim=d.nworld,
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],
outputs=[ctx.beta],
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,
)
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(
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])
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+27 -5
View File
@@ -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
Vendored
+1 -1
View File
@@ -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
+8 -17
View File
@@ -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
+28 -34
View File
@@ -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,