diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 487fd821f..f05aa69c2 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -23,6 +23,7 @@ from .types import vec5 from .types import vec11 from .warp_util import event_scope +from .warp_util import nested_kernel wp.set_module_options({"enable_backward": False}) @@ -54,826 +55,442 @@ def _zero_constraint_counts( nefc_out[worldid] = 0 -@wp.func -def _update_efc_row( - # In: - worldid: int, - timestep: float, - refsafe: int, - efcid: int, - pos_aref: float, - pos_imp: float, - invweight: float, - solref: wp.vec2, - solimp: vec5, - margin: float, - vel: float, - frictionloss: float, - type: int, - id: int, - # Data out: - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - # Calculate kbi - timeconst = solref[0] - dampratio = solref[1] - dmin = solimp[0] - dmax = solimp[1] - width = solimp[2] - mid = solimp[3] - power = solimp[4] - - # TODO(team): wp.static? - if not refsafe: - timeconst = wp.max(timeconst, 2.0 * timestep) - - dmin = wp.clamp(dmin, types.MJ_MINIMP, types.MJ_MAXIMP) - dmax = wp.clamp(dmax, types.MJ_MINIMP, types.MJ_MAXIMP) - width = wp.max(types.MJ_MINVAL, width) - mid = wp.clamp(mid, types.MJ_MINIMP, types.MJ_MAXIMP) - power = wp.max(1.0, power) - - # See https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters - k = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) - b = 2.0 / (dmax * timeconst) - k = wp.where(solref[0] <= 0, -solref[0] / (dmax * dmax), k) - b = wp.where(solref[1] <= 0, -solref[1] / dmax, b) - - imp_x = wp.abs(pos_imp) / width - imp_a = (1.0 / wp.pow(mid, power - 1.0)) * wp.pow(imp_x, power) - imp_b = 1.0 - (1.0 / wp.pow(1.0 - mid, power - 1.0)) * wp.pow(1.0 - imp_x, power) - imp_y = wp.where(imp_x < mid, imp_a, imp_b) - imp = dmin + imp_y * (dmax - dmin) - imp = wp.clamp(imp, dmin, dmax) - imp = wp.where(imp_x > 1.0, dmax, imp) - - # Update constraints - efc_D_out[worldid, efcid] = 1.0 / wp.max(invweight * (1.0 - imp) / imp, types.MJ_MINVAL) - efc_vel_out[worldid, efcid] = vel - efc_aref_out[worldid, efcid] = -k * imp * pos_aref - b * vel - efc_pos_out[worldid, efcid] = pos_aref + margin - efc_margin_out[worldid, efcid] = margin - efc_frictionloss_out[worldid, efcid] = frictionloss - efc_type_out[worldid, efcid] = type - efc_id_out[worldid, efcid] = id - - @wp.kernel -def _efc_equality_connect( - # Model: - nv: int, - nsite: int, - opt_timestep: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), - eq_connect_adr: wp.array(dtype=int), +def _num_equality( # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - njmax_in: int, - # In: - refsafe_in: int, + ne_connect_in: wp.array(dtype=int), + ne_weld_in: wp.array(dtype=int), + ne_jnt_in: wp.array(dtype=int), + ne_ten_in: wp.array(dtype=int), + ne_flex_in: wp.array(dtype=int), # Data out: - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), - ne_connect_out: wp.array(dtype=int), + ne_out: wp.array(dtype=int), ): - """Calculates constraint rows for connect equality constraints.""" - worldid, eqconnectid = wp.tid() - eqid = eq_connect_adr[eqconnectid] - - if not eq_active_in[worldid, eqid]: - return - - wp.atomic_add(ne_connect_out, worldid, 3) - efcid = wp.atomic_add(nefc_out, worldid, 3) - - if efcid + 3 >= njmax_in: - return - - data = eq_data[worldid % eq_data.shape[0], eqid] - anchor1 = wp.vec3f(data[0], data[1], data[2]) - anchor2 = wp.vec3f(data[3], data[4], data[5]) - - obj1id = eq_obj1id[eqid] - obj2id = eq_obj2id[eqid] - - if nsite and eq_objtype[eqid] == types.ObjType.SITE: - # body1id stores the index of site_bodyid. - body1id = site_bodyid[obj1id] - body2id = site_bodyid[obj2id] - pos1 = site_xpos_in[worldid, obj1id] - pos2 = site_xpos_in[worldid, obj2id] - else: - body1id = obj1id - body2id = obj2id - pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor1 - pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor2 - - # error is difference in global positions - pos = pos1 - pos2 - - # compute Jacobian difference (opposite of contact: 0 - 1) - Jqvel = wp.vec3f(0.0, 0.0, 0.0) - for dofid in range(nv): # TODO: parallelize - jacp1, _ = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos1, - body1id, - dofid, - worldid, - ) - jacp2, _ = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos2, - body2id, - dofid, - worldid, - ) - j1mj2 = jacp1 - jacp2 - efc_J_out[worldid, efcid + 0, dofid] = j1mj2[0] - efc_J_out[worldid, efcid + 1, dofid] = j1mj2[1] - efc_J_out[worldid, efcid + 2, dofid] = j1mj2[2] - Jqvel += j1mj2 * qvel_in[worldid, dofid] + worldid = wp.tid() + ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne_flex_in[worldid] + ne_out[worldid] = ne - body_invweight0_id = worldid % body_invweight0.shape[0] - invweight = body_invweight0[body_invweight0_id, body1id][0] + body_invweight0[body_invweight0_id, body2id][0] - pos_imp = wp.length(pos) - solref = eq_solref[worldid % eq_solref.shape[0], eqid] - solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] - timestep = opt_timestep[worldid % opt_timestep.shape[0]] +@event_scope +def make_constraint(m: types.Model, d: types.Data): + """Creates constraint jacobians and other supporting data.""" - for i in range(3): - efcidi = efcid + i + @wp.func + def efc_row( + # In: + worldid: int, + timestep: float, + efcid: int, + pos_aref: float, + pos_imp: float, + invweight: float, + solref: wp.vec2, + solimp: vec5, + margin: float, + vel: float, + frictionloss: float, + type: int, + id: int, + # Out: + type_out: wp.array2d(dtype=int), + id_out: wp.array2d(dtype=int), + pos_out: wp.array2d(dtype=float), + margin_out: wp.array2d(dtype=float), + D_out: wp.array2d(dtype=float), + vel_out: wp.array2d(dtype=float), + aref_out: wp.array2d(dtype=float), + frictionloss_out: wp.array2d(dtype=float), + ): + # calculate kbi + timeconst = solref[0] + dampratio = solref[1] + dmin = solimp[0] + dmax = solimp[1] + width = solimp[2] + mid = solimp[3] + power = solimp[4] + + if wp.static(not (m.opt.disableflags & types.DisableBit.REFSAFE)): + timeconst = wp.max(timeconst, 2.0 * timestep) + + dmin = wp.clamp(dmin, types.MJ_MINIMP, types.MJ_MAXIMP) + dmax = wp.clamp(dmax, types.MJ_MINIMP, types.MJ_MAXIMP) + width = wp.max(types.MJ_MINVAL, width) + mid = wp.clamp(mid, types.MJ_MINIMP, types.MJ_MAXIMP) + power = wp.max(1.0, power) + + # see https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters + dmax_sq = dmax * dmax + k = 1.0 / (dmax_sq * timeconst * timeconst * dampratio * dampratio) + b = 2.0 / (dmax * timeconst) + k = wp.where(solref[0] <= 0, -solref[0] / dmax_sq, k) + b = wp.where(solref[1] <= 0, -solref[1] / dmax, b) + + imp_x = wp.abs(pos_imp) / width + imp_a = (1.0 / wp.pow(mid, power - 1.0)) * wp.pow(imp_x, power) + imp_b = 1.0 - (1.0 / wp.pow(1.0 - mid, power - 1.0)) * wp.pow(1.0 - imp_x, power) + imp_y = wp.where(imp_x < mid, imp_a, imp_b) + imp = dmin + imp_y * (dmax - dmin) + imp = wp.clamp(imp, dmin, dmax) + imp = wp.where(imp_x > 1.0, dmax, imp) + + # set outputs + D_out[worldid, efcid] = 1.0 / wp.max(invweight * (1.0 - imp) / imp, types.MJ_MINVAL) + vel_out[worldid, efcid] = vel + aref_out[worldid, efcid] = -k * imp * pos_aref - b * vel + pos_out[worldid, efcid] = pos_aref + margin + margin_out[worldid, efcid] = margin + frictionloss_out[worldid, efcid] = frictionloss + type_out[worldid, efcid] = type + id_out[worldid, efcid] = id + + @nested_kernel(module="unique", enable_backward=False) + def efc_J_initialize( + # Data out: + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + ): + worldid, efcid, dofid = wp.tid() + + rowadr = efcid * wp.static(m.nv) + if dofid == 0: + efc_J_rownnz_out[worldid, efcid] = wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + + sparseid = rowadr + dofid + efc_J_colind_out[worldid, sparseid] = dofid + efc_J_out[worldid, 0, sparseid] = 0.0 + + @nested_kernel(module="unique", enable_backward=False) + def equality_connect( + # Model: + opt_timestep: wp.array(dtype=float), + body_parentid: wp.array(dtype=int), + body_rootid: wp.array(dtype=int), + body_dofnum: wp.array(dtype=int), + body_dofadr: wp.array(dtype=int), + body_invweight0: wp.array2d(dtype=wp.vec2), + dof_bodyid: wp.array(dtype=int), + dof_parentid: wp.array(dtype=int), + site_bodyid: wp.array(dtype=int), + eq_obj1id: wp.array(dtype=int), + eq_obj2id: wp.array(dtype=int), + eq_objtype: wp.array(dtype=int), + eq_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_data: wp.array2d(dtype=vec11), + eq_connect_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + eq_active_in: wp.array2d(dtype=bool), + xpos_in: wp.array2d(dtype=wp.vec3), + xmat_in: wp.array2d(dtype=wp.mat33), + site_xpos_in: wp.array2d(dtype=wp.vec3), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cdof_in: wp.array2d(dtype=wp.spatial_vector), + # Data out: + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ne_connect_out: wp.array(dtype=int), + ): + """Calculates constraint rows for connect equality constraints.""" + worldid, eqconnectid = wp.tid() + eqid = eq_connect_adr[eqconnectid] + + if not eq_active_in[worldid, eqid]: + return - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcidi, - pos[i], - pos_imp, - invweight, - solref, - solimp, - 0.0, - Jqvel[i], - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + wp.atomic_add(ne_connect_out, worldid, 3) + efcid = wp.atomic_add(nefc_out, worldid, 3) + if efcid >= wp.static(d.njmax - 3): + return -@wp.kernel -def _efc_equality_joint( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - qpos0: wp.array2d(dtype=float), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - dof_invweight0: wp.array2d(dtype=float), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), - eq_jnt_adr: wp.array(dtype=int), - # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), - ne_jnt_out: wp.array(dtype=int), -): - worldid, eqjntid = wp.tid() - eqid = eq_jnt_adr[eqjntid] - - if not eq_active_in[worldid, eqid]: - return - - wp.atomic_add(ne_jnt_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) - - if efcid >= njmax_in: - return - - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - - jntid_1 = eq_obj1id[eqid] - jntid_2 = eq_obj2id[eqid] - data = eq_data[worldid % eq_data.shape[0], eqid] - dofadr1 = jnt_dofadr[jntid_1] - qposadr1 = jnt_qposadr[jntid_1] - efc_J_out[worldid, efcid, dofadr1] = 1.0 - qpos0_id = worldid % qpos0.shape[0] - dof_invweight0_id = worldid % dof_invweight0.shape[0] - - if jntid_2 > -1: - # Two joint constraint - qposadr2 = jnt_qposadr[jntid_2] - dofadr2 = jnt_dofadr[jntid_2] - dif = qpos_in[worldid, qposadr2] - qpos0[worldid, qposadr2] - - # Horner's method for polynomials - rhs = data[0] + dif * (data[1] + dif * (data[2] + dif * (data[3] + dif * data[4]))) - deriv_2 = data[1] + dif * (2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4])) - - pos = qpos_in[worldid, qposadr1] - qpos0[qpos0_id, qposadr1] - rhs - Jqvel = qvel_in[worldid, dofadr1] - qvel_in[worldid, dofadr2] * deriv_2 - invweight = dof_invweight0[dof_invweight0_id, dofadr1] + dof_invweight0[dof_invweight0_id, dofadr2] - - efc_J_out[worldid, efcid, dofadr2] = -deriv_2 - else: - # Single joint constraint - pos = qpos_in[worldid, qposadr1] - qpos0[qpos0_id, qposadr1] - data[0] - Jqvel = qvel_in[worldid, dofadr1] - invweight = dof_invweight0[dof_invweight0_id, dofadr1] - - # Update constraint parameters - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - invweight, - eq_solref[worldid % eq_solref.shape[0], eqid], - eq_solimp[worldid % eq_solimp.shape[0], eqid], - 0.0, - Jqvel, - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + data = eq_data[worldid % eq_data.shape[0], eqid] + anchor1 = wp.vec3f(data[0], data[1], data[2]) + anchor2 = wp.vec3f(data[3], data[4], data[5]) + obj1id = eq_obj1id[eqid] + obj2id = eq_obj2id[eqid] -@wp.kernel -def _efc_equality_tendon( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), - tendon_length0: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), - eq_ten_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - ten_J_in: wp.array3d(dtype=float), - ten_length_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), - ne_ten_out: wp.array(dtype=int), -): - worldid, eqtenid = wp.tid() - eqid = eq_ten_adr[eqtenid] - - if not eq_active_in[worldid, eqid]: - return - - wp.atomic_add(ne_ten_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) - - if efcid >= njmax_in: - return - - obj1id = eq_obj1id[eqid] - obj2id = eq_obj2id[eqid] - - data = eq_data[worldid % eq_data.shape[0], eqid] - solref = eq_solref[worldid % eq_solref.shape[0], eqid] - solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] - tendon_length0_id = worldid % tendon_length0.shape[0] - tendon_invweight0_id = worldid % tendon_invweight0.shape[0] - pos1 = ten_length_in[worldid, obj1id] - tendon_length0[tendon_length0_id, obj1id] - jac1 = ten_J_in[worldid, obj1id] - - if obj2id > -1: - invweight = tendon_invweight0[tendon_invweight0_id, obj1id] + tendon_invweight0[tendon_invweight0_id, obj2id] - - pos2 = ten_length_in[worldid, obj2id] - tendon_length0[tendon_length0_id, obj2id] - jac2 = ten_J_in[worldid, obj2id] - - dif = pos2 - dif2 = dif * dif - dif3 = dif2 * dif - dif4 = dif3 * dif - - pos = pos1 - (data[0] + data[1] * dif + data[2] * dif2 + data[3] * dif3 + data[4] * dif4) - deriv = data[1] + 2.0 * data[2] * dif + 3.0 * data[3] * dif2 + 4.0 * data[4] * dif3 - else: - invweight = tendon_invweight0[tendon_invweight0_id, obj1id] - pos = pos1 - data[0] - deriv = 0.0 - - Jqvel = float(0.0) - for i in range(nv): - if deriv != 0.0: - J = jac1[i] + jac2[i] * -deriv + if wp.static(m.nsite) and eq_objtype[eqid] == types.ObjType.SITE: + body1 = site_bodyid[obj1id] + body2 = site_bodyid[obj2id] + pos1 = site_xpos_in[worldid, obj1id] + pos2 = site_xpos_in[worldid, obj2id] else: - J = jac1[i] - efc_J_out[worldid, efcid, i] = J - Jqvel += J * qvel_in[worldid, i] - - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - invweight, - solref, - solimp, - 0.0, - Jqvel, - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + body1 = obj1id + body2 = obj2id + pos1 = xpos_in[worldid, body1] + xmat_in[worldid, body1] @ anchor1 + pos2 = xpos_in[worldid, body2] + xmat_in[worldid, body2] @ anchor2 + + # error is difference in global positions + pos = pos1 - pos2 + + # compute Jacobian difference (opposite of contact: 0 - 1) + Jqvel = wp.vec3f(0.0, 0.0, 0.0) + + if wp.static(m.opt.is_sparse): + while body1 > 0 and body_dofnum[body1] == 0: + body1 = body_parentid[body1] + while body2 > 0 and body_dofnum[body2] == 0: + body2 = body_parentid[body2] + + da1 = int(body_dofadr[body1] + body_dofnum[body1] - 1) + da2 = int(body_dofadr[body2] + body_dofnum[body2] - 1) + + efcid0 = efcid + 0 + efcid1 = efcid + 1 + efcid2 = efcid + 2 + + rowadr0 = efcid0 * wp.static(m.nv) + rowadr1 = efcid1 * wp.static(m.nv) + rowadr2 = efcid2 * wp.static(m.nv) + + efc_J_rowadr_out[worldid, efcid0] = rowadr0 + efc_J_rowadr_out[worldid, efcid1] = rowadr1 + efc_J_rowadr_out[worldid, efcid2] = rowadr2 + + rownnz = int(0) + + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + if da1 == da: + da1 = dof_parentid[da1] + if da2 == da: + da2 = dof_parentid[da2] + + jacp1, _ = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + da, + worldid, + ) + jacp2, _ = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + da, + worldid, + ) + j1mj2 = jacp1 - jacp2 + sparseid0 = rowadr0 + rownnz + sparseid1 = rowadr1 + rownnz + sparseid2 = rowadr2 + rownnz -@wp.kernel -def _efc_equality_flex( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - flexedge_length0: wp.array(dtype=float), - flexedge_invweight0: wp.array(dtype=float), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_flex_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - flexedge_J_in: wp.array3d(dtype=float), - flexedge_length_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), - ne_flex_out: wp.array(dtype=int), -): - worldid, eqflexid, edgeid = wp.tid() - eqid = eq_flex_adr[eqflexid] - - wp.atomic_add(ne_flex_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) - - if efcid >= njmax_in: - return - - pos = flexedge_length_in[worldid, edgeid] - flexedge_length0[edgeid] - solref = eq_solref[worldid % eq_solref.shape[0], eqid] - solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] - - Jqvel = float(0.0) - for i in range(nv): - J = flexedge_J_in[worldid, edgeid, i] - efc_J_out[worldid, efcid, i] = J - Jqvel += J * qvel_in[worldid, i] - - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - pos, - pos, - flexedge_invweight0[edgeid], - solref, - solimp, - 0.0, - Jqvel, - 0.0, - ConstraintType.EQUALITY, - eqid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + efc_J_colind_out[worldid, sparseid0] = da + efc_J_colind_out[worldid, sparseid1] = da + efc_J_colind_out[worldid, sparseid2] = da + efc_J_out[worldid, 0, sparseid0] = j1mj2[0] + efc_J_out[worldid, 0, sparseid1] = j1mj2[1] + efc_J_out[worldid, 0, sparseid2] = j1mj2[2] -@wp.kernel -def _efc_friction_dof( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - dof_solref: wp.array2d(dtype=wp.vec2), - dof_solimp: wp.array2d(dtype=vec5), - dof_frictionloss: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), - # Data in: - qvel_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nf_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - worldid, dofid = wp.tid() - - dof_frictionloss_id = worldid % dof_frictionloss.shape[0] - - if dof_frictionloss[dof_frictionloss_id, dofid] <= 0.0: - return - - wp.atomic_add(nf_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) - - if efcid >= njmax_in: - return - - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - - efc_J_out[worldid, efcid, dofid] = 1.0 - Jqvel = qvel_in[worldid, dofid] - - dof_invweight0_id = worldid % dof_invweight0.shape[0] - dof_solref_id = worldid % dof_solref.shape[0] - dof_solimp_id = worldid % dof_solimp.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - 0.0, - 0.0, - dof_invweight0[dof_invweight0_id, dofid], - dof_solref[dof_solref_id, dofid], - dof_solimp[dof_solimp_id, dofid], - 0.0, - Jqvel, - dof_frictionloss[dof_frictionloss_id, dofid], - ConstraintType.FRICTION_DOF, - dofid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + Jqvel += j1mj2 * qvel_in[worldid, da] + rownnz += 1 -@wp.kernel -def _efc_friction_tendon( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - tendon_solref_fri: wp.array2d(dtype=wp.vec2), - tendon_solimp_fri: wp.array2d(dtype=vec5), - tendon_frictionloss: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), - # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array3d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nf_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - worldid, tenid = wp.tid() - - tendon_frictionloss_id = worldid % tendon_frictionloss.shape[0] - - frictionloss = tendon_frictionloss[tendon_frictionloss_id, tenid] - if frictionloss <= 0.0: - return - - wp.atomic_add(nf_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) - - if efcid >= njmax_in: - return - - Jqvel = float(0.0) - - # TODO(team): parallelize - for i in range(nv): - J = ten_J_in[worldid, tenid, i] - efc_J_out[worldid, efcid, i] = J - Jqvel += J * qvel_in[worldid, i] - - tendon_invweight0_id = worldid % tendon_invweight0.shape[0] - tendon_solref_fri_id = worldid % tendon_solref_fri.shape[0] - tendon_solimp_fri_id = worldid % tendon_solimp_fri.shape[0] - _update_efc_row( - worldid, - opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, - efcid, - 0.0, - 0.0, - tendon_invweight0[tendon_invweight0_id, tenid], - tendon_solref_fri[tendon_solref_fri_id, tenid], - tendon_solimp_fri[tendon_solimp_fri_id, tenid], - 0.0, - Jqvel, - frictionloss, - ConstraintType.FRICTION_TENDON, - tenid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + efc_J_rownnz_out[worldid, efcid0] = rownnz + efc_J_rownnz_out[worldid, efcid1] = rownnz + efc_J_rownnz_out[worldid, efcid2] = rownnz + else: + # TODO(team): dof tree traversal + for dofid in range(wp.static(m.nv)): + jacp1, _ = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + dofid, + worldid, + ) + jacp2, _ = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + dofid, + worldid, + ) + j1mj2 = jacp1 - jacp2 + efc_J_out[worldid, efcid + 0, dofid] = j1mj2[0] + efc_J_out[worldid, efcid + 1, dofid] = j1mj2[1] + efc_J_out[worldid, efcid + 2, dofid] = j1mj2[2] -@wp.kernel -def _efc_equality_weld( - # Model: - nv: int, - nsite: int, - opt_timestep: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - site_bodyid: wp.array(dtype=int), - site_quat: wp.array2d(dtype=wp.quat), - eq_obj1id: wp.array(dtype=int), - eq_obj2id: wp.array(dtype=int), - eq_objtype: wp.array(dtype=int), - eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array2d(dtype=vec11), - eq_wld_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - eq_active_in: wp.array2d(dtype=bool), - xpos_in: wp.array2d(dtype=wp.vec3), - xquat_in: wp.array2d(dtype=wp.quat), - xmat_in: wp.array2d(dtype=wp.mat33), - site_xpos_in: wp.array2d(dtype=wp.vec3), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), - ne_weld_out: wp.array(dtype=int), -): - worldid, eqweldid = wp.tid() - eqid = eq_wld_adr[eqweldid] - - if not eq_active_in[worldid, eqid]: - return - - wp.atomic_add(ne_weld_out, worldid, 6) - efcid = wp.atomic_add(nefc_out, worldid, 6) - - if efcid + 6 >= njmax_in: - return - - is_site = eq_objtype[eqid] == types.ObjType.SITE and nsite > 0 - - obj1id = eq_obj1id[eqid] - obj2id = eq_obj2id[eqid] - - data = eq_data[worldid % eq_data.shape[0], eqid] - anchor1 = wp.vec3(data[0], data[1], data[2]) - anchor2 = wp.vec3(data[3], data[4], data[5]) - relpose = wp.quat(data[6], data[7], data[8], data[9]) - torquescale = data[10] - - if is_site: - # body1id stores the index of site_bodyid. - body1id = site_bodyid[obj1id] - body2id = site_bodyid[obj2id] - pos1 = site_xpos_in[worldid, obj1id] - pos2 = site_xpos_in[worldid, obj2id] - - site_quat_id = worldid % site_quat.shape[0] - quat = math.mul_quat(xquat_in[worldid, body1id], site_quat[site_quat_id, obj1id]) - quat1 = math.quat_inv(math.mul_quat(xquat_in[worldid, body2id], site_quat[site_quat_id, obj2id])) - - else: - body1id = obj1id - body2id = obj2id - pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor2 - pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor1 - - quat = math.mul_quat(xquat_in[worldid, body1id], relpose) - quat1 = math.quat_inv(xquat_in[worldid, body2id]) - - # compute Jacobian difference (opposite of contact: 0 - 1) - Jqvelp = wp.vec3f(0.0, 0.0, 0.0) - Jqvelr = wp.vec3f(0.0, 0.0, 0.0) - - for dofid in range(nv): # TODO: parallelize - jacp1, jacr1 = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos1, - body1id, - dofid, - worldid, - ) - jacp2, jacr2 = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - pos2, - body2id, - dofid, - worldid, - ) + Jqvel += j1mj2 * qvel_in[worldid, dofid] - jacdifp = jacp1 - jacp2 - for i in range(3): - efc_J_out[worldid, efcid + i, dofid] = jacdifp[i] + body_invweight0_id = worldid % body_invweight0.shape[0] + invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] + pos_imp = wp.length(pos) - jacdifr = (jacr1 - jacr2) * torquescale - jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat) - jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3]) + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + timestep = opt_timestep[worldid % opt_timestep.shape[0]] for i in range(3): - efc_J_out[worldid, efcid + 3 + i, dofid] = jacdifr[i] + efcidi = efcid + i - Jqvelp += jacdifp * qvel_in[worldid, dofid] - Jqvelr += jacdifr * qvel_in[worldid, dofid] - - # error is difference in global position and orientation - cpos = pos1 - pos2 + efc_row( + worldid, + timestep, + efcidi, + pos[i], + pos_imp, + invweight, + solref, + solimp, + 0.0, + Jqvel[i], + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) - crotq = math.mul_quat(quat1, quat) # copy axis components - crot = wp.vec3(crotq[1], crotq[2], crotq[3]) * torquescale + @nested_kernel(module="unique", enable_backward=False) + def equality_joint( + # Model: + opt_timestep: wp.array(dtype=float), + qpos0: wp.array2d(dtype=float), + jnt_qposadr: wp.array(dtype=int), + jnt_dofadr: wp.array(dtype=int), + dof_invweight0: wp.array2d(dtype=float), + eq_obj1id: wp.array(dtype=int), + eq_obj2id: wp.array(dtype=int), + eq_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_data: wp.array2d(dtype=vec11), + eq_jnt_adr: wp.array(dtype=int), + # Data in: + qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), + eq_active_in: wp.array2d(dtype=bool), + # Data out: + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ne_jnt_out: wp.array(dtype=int), + ): + worldid, eqjntid = wp.tid() + eqid = eq_jnt_adr[eqjntid] + + if not eq_active_in[worldid, eqid]: + return - body_invweight0_id = worldid % body_invweight0.shape[0] - invweight_t = body_invweight0[body_invweight0_id, body1id][0] + body_invweight0[body_invweight0_id, body2id][0] + wp.atomic_add(ne_jnt_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) - pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) + if efcid >= wp.static(d.njmax): + return - solref = eq_solref[worldid % eq_solref.shape[0], eqid] - solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + jntid_1 = eq_obj1id[eqid] + jntid_2 = eq_obj2id[eqid] + data = eq_data[worldid % eq_data.shape[0], eqid] + dofadr1 = jnt_dofadr[jntid_1] + qposadr1 = jnt_qposadr[jntid_1] + qpos0_id = worldid % qpos0.shape[0] + dof_invweight0_id = worldid % dof_invweight0.shape[0] - timestep = opt_timestep[worldid % opt_timestep.shape[0]] + if wp.static(m.opt.is_sparse): + if jntid_2 > -1: + rownnz = 2 + else: + rownnz = 1 + efc_J_rownnz_out[worldid, efcid] = rownnz + rowadr = efcid * wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + efc_J_colind_out[worldid, rowadr] = dofadr1 + efc_J_out[worldid, 0, rowadr] = 1.0 + else: + for i in range(wp.static(m.nv)): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dofadr1] = 1.0 + + if jntid_2 > -1: + # Two joint constraint + qposadr2 = jnt_qposadr[jntid_2] + dofadr2 = jnt_dofadr[jntid_2] + dif = qpos_in[worldid, qposadr2] - qpos0[worldid, qposadr2] + + # Horner's method for polynomials + rhs = data[0] + dif * (data[1] + dif * (data[2] + dif * (data[3] + dif * data[4]))) + deriv_2 = data[1] + dif * (2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4])) + + pos = qpos_in[worldid, qposadr1] - qpos0[qpos0_id, qposadr1] - rhs + Jqvel = qvel_in[worldid, dofadr1] - qvel_in[worldid, dofadr2] * deriv_2 + invweight = dof_invweight0[dof_invweight0_id, dofadr1] + dof_invweight0[dof_invweight0_id, dofadr2] + + if wp.static(m.opt.is_sparse): + sparseid = rowadr + 1 + efc_J_colind_out[worldid, sparseid] = dofadr2 + efc_J_out[worldid, 0, sparseid] = -deriv_2 + else: + efc_J_out[worldid, efcid, dofadr2] = -deriv_2 + else: + # Single joint constraint + pos = qpos_in[worldid, qposadr1] - qpos0[qpos0_id, qposadr1] - data[0] + Jqvel = qvel_in[worldid, dofadr1] + invweight = dof_invweight0[dof_invweight0_id, dofadr1] - for i in range(3): - _update_efc_row( + # Update constraint parameters + efc_row( worldid, - timestep, - refsafe_in, - efcid + i, - cpos[i], - pos_imp, - invweight_t, - solref, - solimp, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + invweight, + eq_solref[worldid % eq_solref.shape[0], eqid], + eq_solimp[worldid % eq_solimp.shape[0], eqid], 0.0, - Jqvelp[i], + Jqvel, 0.0, ConstraintType.EQUALITY, eqid, @@ -887,21 +504,102 @@ def _efc_equality_weld( efc_frictionloss_out, ) - invweight_r = body_invweight0[body_invweight0_id, body1id][1] + body_invweight0[body_invweight0_id, body2id][1] + @nested_kernel(module="unique", enable_backward=False) + def equality_tendon( + # Model: + opt_timestep: wp.array(dtype=float), + eq_obj1id: wp.array(dtype=int), + eq_obj2id: wp.array(dtype=int), + eq_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_data: wp.array2d(dtype=vec11), + tendon_length0: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), + eq_ten_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + eq_active_in: wp.array2d(dtype=bool), + ten_J_in: wp.array3d(dtype=float), + ten_length_in: wp.array2d(dtype=float), + # Data out: + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ne_ten_out: wp.array(dtype=int), + ): + worldid, eqtenid = wp.tid() + eqid = eq_ten_adr[eqtenid] + + if not eq_active_in[worldid, eqid]: + return + + wp.atomic_add(ne_ten_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= wp.static(d.njmax): + return + + obj1id = eq_obj1id[eqid] + obj2id = eq_obj2id[eqid] + + data = eq_data[worldid % eq_data.shape[0], eqid] + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + tendon_length0_id = worldid % tendon_length0.shape[0] + tendon_invweight0_id = worldid % tendon_invweight0.shape[0] + pos1 = ten_length_in[worldid, obj1id] - tendon_length0[tendon_length0_id, obj1id] + jac1 = ten_J_in[worldid, obj1id] + + if obj2id > -1: + invweight = tendon_invweight0[tendon_invweight0_id, obj1id] + tendon_invweight0[tendon_invweight0_id, obj2id] + + pos2 = ten_length_in[worldid, obj2id] - tendon_length0[tendon_length0_id, obj2id] + jac2 = ten_J_in[worldid, obj2id] + + dif = pos2 + dif2 = dif * dif + dif3 = dif2 * dif + dif4 = dif3 * dif + + pos = pos1 - (data[0] + data[1] * dif + data[2] * dif2 + data[3] * dif3 + data[4] * dif4) + deriv = data[1] + 2.0 * data[2] * dif + 3.0 * data[3] * dif2 + 4.0 * data[4] * dif3 + else: + invweight = tendon_invweight0[tendon_invweight0_id, obj1id] + pos = pos1 - data[0] + deriv = 0.0 + + Jqvel = float(0.0) + + # TODO(team): efc_J_rownnz, efc_J_rowadr, efc_J_colind + for i in range(wp.static(m.nv)): + if deriv != 0.0: + J = jac1[i] + jac2[i] * -deriv + else: + J = jac1[i] + if wp.static(m.opt.is_sparse): + efc_J_out[worldid, 0, efcid * wp.static(m.nv) + i] = J + else: + efc_J_out[worldid, efcid, i] = J + Jqvel += J * qvel_in[worldid, i] - for i in range(3): - _update_efc_row( + efc_row( worldid, - timestep, - refsafe_in, - efcid + 3 + i, - crot[i], - pos_imp, - invweight_r, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + invweight, solref, solimp, 0.0, - Jqvelr[i], + Jqvel, 0.0, ConstraintType.EQUALITY, eqid, @@ -915,85 +613,71 @@ def _efc_equality_weld( efc_frictionloss_out, ) - -@wp.kernel -def _efc_limit_slide_hinge( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array2d(dtype=vec5), - jnt_range: wp.array2d(dtype=wp.vec2), - jnt_margin: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), - jnt_limited_slide_hinge_adr: wp.array(dtype=int), - # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - worldid, jntlimitedid = wp.tid() - jntid = jnt_limited_slide_hinge_adr[jntlimitedid] - jnt_range_id = worldid % jnt_range.shape[0] - jntrange = jnt_range[jnt_range_id, jntid] - - qpos = qpos_in[worldid, jnt_qposadr[jntid]] - jnt_margin_id = worldid % jnt_margin.shape[0] - jntmargin = jnt_margin[jnt_margin_id, jntid] - dist_min, dist_max = qpos - jntrange[0], jntrange[1] - qpos - pos = wp.min(dist_min, dist_max) - jntmargin - active = pos < 0 - - if active: - wp.atomic_add(nl_out, worldid, 1) + @nested_kernel(module="unique", enable_backward=False) + def equality_flex( + # Model: + opt_timestep: wp.array(dtype=float), + flexedge_length0: wp.array(dtype=float), + flexedge_invweight0: wp.array(dtype=float), + eq_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_flex_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + flexedge_J_in: wp.array3d(dtype=float), + flexedge_length_in: wp.array2d(dtype=float), + # Data out: + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ne_flex_out: wp.array(dtype=int), + ): + worldid, eqflexid, edgeid = wp.tid() + eqid = eq_flex_adr[eqflexid] + + wp.atomic_add(ne_flex_out, worldid, 1) efcid = wp.atomic_add(nefc_out, worldid, 1) - if efcid >= njmax_in: + if efcid >= wp.static(d.njmax): return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 + pos = flexedge_length_in[worldid, edgeid] - flexedge_length0[edgeid] + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] - dofadr = jnt_dofadr[jntid] + Jqvel = float(0.0) - J = float(dist_min < dist_max) * 2.0 - 1.0 - efc_J_out[worldid, efcid, dofadr] = J - Jqvel = J * qvel_in[worldid, dofadr] + # TODO(team): efc.J_rownnz, efc.J_rowadr, efc.J_colind + for i in range(wp.static(m.nv)): + # TODO(team): sparse flexedge_J + J = flexedge_J_in[worldid, edgeid, i] + if wp.static(m.opt.is_sparse): + efc_J_out[worldid, 0, efcid * wp.static(m.nv) + i] = J + else: + efc_J_out[worldid, efcid, i] = J + Jqvel += J * qvel_in[worldid, i] - dof_invweight0_id = worldid % dof_invweight0.shape[0] - jnt_solref_id = worldid % jnt_solref.shape[0] - jnt_solimp_id = worldid % jnt_solimp.shape[0] - _update_efc_row( + efc_row( worldid, opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, efcid, pos, pos, - dof_invweight0[dof_invweight0_id, dofadr], - jnt_solref[jnt_solref_id, jntid], - jnt_solimp[jnt_solimp_id, jntid], - jntmargin, + flexedge_invweight0[edgeid], + solref, + solimp, + 0.0, Jqvel, 0.0, - ConstraintType.LIMIT_JOINT, - jntid, + ConstraintType.EQUALITY, + eqid, efc_type_out, efc_id_out, efc_pos_out, @@ -1004,94 +688,374 @@ def _efc_limit_slide_hinge( efc_frictionloss_out, ) + @nested_kernel(module="unique", enable_backward=False) + def equality_weld( + # Model: + opt_timestep: wp.array(dtype=float), + body_parentid: wp.array(dtype=int), + body_rootid: wp.array(dtype=int), + body_dofnum: wp.array(dtype=int), + body_dofadr: wp.array(dtype=int), + body_invweight0: wp.array2d(dtype=wp.vec2), + dof_bodyid: wp.array(dtype=int), + dof_parentid: wp.array(dtype=int), + site_bodyid: wp.array(dtype=int), + site_quat: wp.array2d(dtype=wp.quat), + eq_obj1id: wp.array(dtype=int), + eq_obj2id: wp.array(dtype=int), + eq_objtype: wp.array(dtype=int), + eq_solref: wp.array2d(dtype=wp.vec2), + eq_solimp: wp.array2d(dtype=vec5), + eq_data: wp.array2d(dtype=vec11), + eq_wld_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + eq_active_in: wp.array2d(dtype=bool), + xpos_in: wp.array2d(dtype=wp.vec3), + xquat_in: wp.array2d(dtype=wp.quat), + xmat_in: wp.array2d(dtype=wp.mat33), + site_xpos_in: wp.array2d(dtype=wp.vec3), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cdof_in: wp.array2d(dtype=wp.spatial_vector), + # Data out: + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ne_weld_out: wp.array(dtype=int), + ): + worldid, eqweldid = wp.tid() + eqid = eq_wld_adr[eqweldid] + + if not eq_active_in[worldid, eqid]: + return -@wp.kernel -def _efc_limit_ball( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - jnt_qposadr: wp.array(dtype=int), - jnt_dofadr: wp.array(dtype=int), - jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array2d(dtype=vec5), - jnt_range: wp.array2d(dtype=wp.vec2), - jnt_margin: wp.array2d(dtype=float), - dof_invweight0: wp.array2d(dtype=float), - jnt_limited_ball_adr: wp.array(dtype=int), - # Data in: - qpos_in: wp.array2d(dtype=float), - qvel_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - worldid, jntlimitedid = wp.tid() - jntid = jnt_limited_ball_adr[jntlimitedid] - qposadr = jnt_qposadr[jntid] - - qpos = qpos_in[worldid] - jnt_quat = wp.quat(qpos[qposadr + 0], qpos[qposadr + 1], qpos[qposadr + 2], qpos[qposadr + 3]) - jnt_quat = wp.normalize(jnt_quat) - axis_angle = math.quat_to_vel(jnt_quat) - jnt_range_id = worldid % jnt_range.shape[0] - jntrange = jnt_range[jnt_range_id, jntid] - axis, angle = math.normalize_with_norm(axis_angle) - jnt_margin_id = worldid % jnt_margin.shape[0] - jntmargin = jnt_margin[jnt_margin_id, jntid] - - pos = wp.max(jntrange[0], jntrange[1]) - angle - jntmargin - active = pos < 0 - - if active: - wp.atomic_add(nl_out, worldid, 1) - efcid = wp.atomic_add(nefc_out, worldid, 1) + wp.atomic_add(ne_weld_out, worldid, 6) + efcid = wp.atomic_add(nefc_out, worldid, 6) + + if efcid >= wp.static(d.njmax - 6): + return + + is_site = eq_objtype[eqid] == types.ObjType.SITE and wp.static(m.nsite > 0) + + obj1id = eq_obj1id[eqid] + obj2id = eq_obj2id[eqid] + + data = eq_data[worldid % eq_data.shape[0], eqid] + anchor1 = wp.vec3(data[0], data[1], data[2]) + anchor2 = wp.vec3(data[3], data[4], data[5]) + relpose = wp.quat(data[6], data[7], data[8], data[9]) + torquescale = data[10] + + if is_site: + body1 = site_bodyid[obj1id] + body2 = site_bodyid[obj2id] + pos1 = site_xpos_in[worldid, obj1id] + pos2 = site_xpos_in[worldid, obj2id] + + site_quat_id = worldid % site_quat.shape[0] + quat = math.mul_quat(xquat_in[worldid, body1], site_quat[site_quat_id, obj1id]) + quat1 = math.quat_inv(math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])) + + else: + body1 = obj1id + body2 = obj2id + pos1 = xpos_in[worldid, body1] + xmat_in[worldid, body1] @ anchor2 + pos2 = xpos_in[worldid, body2] + xmat_in[worldid, body2] @ anchor1 + + quat = math.mul_quat(xquat_in[worldid, body1], relpose) + quat1 = math.quat_inv(xquat_in[worldid, body2]) + + # compute Jacobian difference (opposite of contact: 0 - 1) + Jqvelp = wp.vec3f(0.0, 0.0, 0.0) + Jqvelr = wp.vec3f(0.0, 0.0, 0.0) + + if wp.static(m.opt.is_sparse): + while body1 > 0 and body_dofnum[body1] == 0: + body1 = body_parentid[body1] + while body2 > 0 and body_dofnum[body2] == 0: + body2 = body_parentid[body2] + + da1 = int(body_dofadr[body1] + body_dofnum[body1] - 1) + da2 = int(body_dofadr[body2] + body_dofnum[body2] - 1) + + efcid0 = efcid + 0 + efcid1 = efcid + 1 + efcid2 = efcid + 2 + efcid3 = efcid + 3 + efcid4 = efcid + 4 + efcid5 = efcid + 5 + + rowadr0 = efcid0 * wp.static(m.nv) + rowadr1 = efcid1 * wp.static(m.nv) + rowadr2 = efcid2 * wp.static(m.nv) + rowadr3 = efcid3 * wp.static(m.nv) + rowadr4 = efcid4 * wp.static(m.nv) + rowadr5 = efcid5 * wp.static(m.nv) + + efc_J_rowadr_out[worldid, efcid0] = rowadr0 + efc_J_rowadr_out[worldid, efcid1] = rowadr1 + efc_J_rowadr_out[worldid, efcid2] = rowadr2 + efc_J_rowadr_out[worldid, efcid3] = rowadr3 + efc_J_rowadr_out[worldid, efcid4] = rowadr4 + efc_J_rowadr_out[worldid, efcid5] = rowadr5 + + rownnz = int(0) + + while da1 >= 0 or da2 >= 0: + da = wp.max(da1, da2) + if da1 == da: + da1 = dof_parentid[da] + if da2 == da: + da2 = dof_parentid[da] + + jacp1, jacr1 = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + da, + worldid, + ) + jacp2, jacr2 = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + da, + worldid, + ) + + jacdifp = jacp1 - jacp2 + + jacdifr = (jacr1 - jacr2) * torquescale + jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat) + jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3]) + + sparseid0 = rowadr0 + rownnz + sparseid1 = rowadr1 + rownnz + sparseid2 = rowadr2 + rownnz + sparseid3 = rowadr3 + rownnz + sparseid4 = rowadr4 + rownnz + sparseid5 = rowadr5 + rownnz + + efc_J_colind_out[worldid, sparseid0] = da + efc_J_colind_out[worldid, sparseid1] = da + efc_J_colind_out[worldid, sparseid2] = da + efc_J_colind_out[worldid, sparseid3] = da + efc_J_colind_out[worldid, sparseid4] = da + efc_J_colind_out[worldid, sparseid5] = da + + efc_J_out[worldid, 0, sparseid0] = jacdifp[0] + efc_J_out[worldid, 0, sparseid1] = jacdifp[1] + efc_J_out[worldid, 0, sparseid2] = jacdifp[2] + efc_J_out[worldid, 0, sparseid3] = jacdifr[0] + efc_J_out[worldid, 0, sparseid4] = jacdifr[1] + efc_J_out[worldid, 0, sparseid5] = jacdifr[2] + + Jqvelp += jacdifp * qvel_in[worldid, da] + Jqvelr += jacdifr * qvel_in[worldid, da] + + rownnz += 1 + + efc_J_rownnz_out[worldid, efcid0] = rownnz + efc_J_rownnz_out[worldid, efcid1] = rownnz + efc_J_rownnz_out[worldid, efcid2] = rownnz + efc_J_rownnz_out[worldid, efcid3] = rownnz + efc_J_rownnz_out[worldid, efcid4] = rownnz + efc_J_rownnz_out[worldid, efcid5] = rownnz + else: + for dofid in range(wp.static(m.nv)): + jacp1, jacr1 = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos1, + body1, + dofid, + worldid, + ) + jacp2, jacr2 = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + pos2, + body2, + dofid, + worldid, + ) + + jacdifp = jacp1 - jacp2 + + for i in range(3): + efc_J_out[worldid, efcid + i, dofid] = jacdifp[i] + + jacdifr = (jacr1 - jacr2) * torquescale + jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat) + jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3]) + + for i in range(3): + efc_J_out[worldid, efcid + 3 + i, dofid] = jacdifr[i] + + Jqvelp += jacdifp * qvel_in[worldid, dofid] + Jqvelr += jacdifr * qvel_in[worldid, dofid] + + # error is difference in global position and orientation + cpos = pos1 - pos2 + + crotq = math.mul_quat(quat1, quat) # copy axis components + crot = wp.vec3(crotq[1], crotq[2], crotq[3]) * torquescale + + body_invweight0_id = worldid % body_invweight0.shape[0] + invweight_t = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] - if efcid >= njmax_in: + pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) + + solref = eq_solref[worldid % eq_solref.shape[0], eqid] + solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid] + + timestep = opt_timestep[worldid % opt_timestep.shape[0]] + + for i in range(3): + efc_row( + worldid, + timestep, + efcid + i, + cpos[i], + pos_imp, + invweight_t, + solref, + solimp, + 0.0, + Jqvelp[i], + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) + + invweight_r = body_invweight0[body_invweight0_id, body1][1] + body_invweight0[body_invweight0_id, body2][1] + + for i in range(3): + efc_row( + worldid, + timestep, + efcid + 3 + i, + crot[i], + pos_imp, + invweight_r, + solref, + solimp, + 0.0, + Jqvelr[i], + 0.0, + ConstraintType.EQUALITY, + eqid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) + + @nested_kernel(module="unique", enable_backward=False) + def friction_dof( + # Model: + opt_timestep: wp.array(dtype=float), + dof_solref: wp.array2d(dtype=wp.vec2), + dof_solimp: wp.array2d(dtype=vec5), + dof_frictionloss: wp.array2d(dtype=float), + dof_invweight0: wp.array2d(dtype=float), + # Data in: + qvel_in: wp.array2d(dtype=float), + # Data out: + nf_out: wp.array(dtype=int), + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + worldid, dofid = wp.tid() + + dof_frictionloss_id = worldid % dof_frictionloss.shape[0] + + if dof_frictionloss[dof_frictionloss_id, dofid] <= 0.0: return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 + wp.atomic_add(nf_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) - dofadr = jnt_dofadr[jntid] + if efcid >= wp.static(d.njmax): + return - efc_J_out[worldid, efcid, dofadr + 0] = -axis[0] - efc_J_out[worldid, efcid, dofadr + 1] = -axis[1] - efc_J_out[worldid, efcid, dofadr + 2] = -axis[2] + if wp.static(m.opt.is_sparse): + efc_J_rownnz_out[worldid, efcid] = 1 + rowadr = efcid * wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + efc_J_colind_out[worldid, rowadr] = dofid + efc_J_out[worldid, 0, rowadr] = 1.0 + else: + for i in range(wp.static(m.nv)): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dofid] = 1.0 - Jqvel = -axis[0] * qvel_in[worldid, dofadr + 0] - Jqvel -= axis[1] * qvel_in[worldid, dofadr + 1] - Jqvel -= axis[2] * qvel_in[worldid, dofadr + 2] + Jqvel = qvel_in[worldid, dofid] dof_invweight0_id = worldid % dof_invweight0.shape[0] - jnt_solref_id = worldid % jnt_solref.shape[0] - jnt_solimp_id = worldid % jnt_solimp.shape[0] - _update_efc_row( + dof_solref_id = worldid % dof_solref.shape[0] + dof_solimp_id = worldid % dof_solimp.shape[0] + efc_row( worldid, opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, efcid, - pos, - pos, - dof_invweight0[dof_invweight0_id, dofadr], - jnt_solref[jnt_solref_id, jntid], - jnt_solimp[jnt_solimp_id, jntid], - jntmargin, - Jqvel, 0.0, - ConstraintType.LIMIT_JOINT, - jntid, + 0.0, + dof_invweight0[dof_invweight0_id, dofid], + dof_solref[dof_solref_id, dofid], + dof_solimp[dof_solimp_id, dofid], + 0.0, + Jqvel, + dof_frictionloss[dof_frictionloss_id, dofid], + ConstraintType.FRICTION_DOF, + dofid, efc_type_out, efc_id_out, efc_pos_out, @@ -1102,99 +1066,73 @@ def _efc_limit_ball( efc_frictionloss_out, ) + @nested_kernel(module="unique", enable_backward=False) + def friction_tendon( + # Model: + opt_timestep: wp.array(dtype=float), + tendon_solref_fri: wp.array2d(dtype=wp.vec2), + tendon_solimp_fri: wp.array2d(dtype=vec5), + tendon_frictionloss: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), + # Data in: + qvel_in: wp.array2d(dtype=float), + ten_J_in: wp.array3d(dtype=float), + # Data out: + nf_out: wp.array(dtype=int), + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + worldid, tenid = wp.tid() + + tendon_frictionloss_id = worldid % tendon_frictionloss.shape[0] + + frictionloss = tendon_frictionloss[tendon_frictionloss_id, tenid] + if frictionloss <= 0.0: + return -@wp.kernel -def _efc_limit_tendon( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - jnt_dofadr: wp.array(dtype=int), - tendon_adr: wp.array(dtype=int), - tendon_num: wp.array(dtype=int), - tendon_solref_lim: wp.array2d(dtype=wp.vec2), - tendon_solimp_lim: wp.array2d(dtype=vec5), - tendon_range: wp.array2d(dtype=wp.vec2), - tendon_margin: wp.array2d(dtype=float), - tendon_invweight0: wp.array2d(dtype=float), - wrap_type: wp.array(dtype=int), - wrap_objid: wp.array(dtype=int), - tendon_limited_adr: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - ten_J_in: wp.array3d(dtype=float), - ten_length_in: wp.array2d(dtype=float), - njmax_in: int, - # In: - refsafe_in: int, - # Data out: - nl_out: wp.array(dtype=int), - nefc_out: wp.array(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - worldid, tenlimitedid = wp.tid() - tenid = tendon_limited_adr[tenlimitedid] - - tendon_range_id = worldid % tendon_range.shape[0] - tenrange = tendon_range[tendon_range_id, tenid] - length = ten_length_in[worldid, tenid] - dist_min, dist_max = length - tenrange[0], tenrange[1] - length - tendon_margin_id = worldid % tendon_margin.shape[0] - tenmargin = tendon_margin[tendon_margin_id, tenid] - pos = wp.min(dist_min, dist_max) - tenmargin - active = pos < 0 - - if active: - wp.atomic_add(nl_out, worldid, 1) + wp.atomic_add(nf_out, worldid, 1) efcid = wp.atomic_add(nefc_out, worldid, 1) - if efcid >= njmax_in: + if efcid >= wp.static(d.njmax): return - for i in range(nv): - efc_J_out[worldid, efcid, i] = 0.0 - Jqvel = float(0.0) - scl = float(dist_min < dist_max) * 2.0 - 1.0 - - adr = tendon_adr[tenid] - if wrap_type[adr] == types.WrapType.JOINT: - ten_num = tendon_num[tenid] - for i in range(ten_num): - dofadr = jnt_dofadr[wrap_objid[adr + i]] - J = scl * ten_J_in[worldid, tenid, dofadr] - efc_J_out[worldid, efcid, dofadr] = J - Jqvel += J * qvel_in[worldid, dofadr] - else: - for i in range(nv): - J = scl * ten_J_in[worldid, tenid, i] + + # TODO(team): efc.J_rownnz, efc.J_rowadr, efc.J_colind + for i in range(wp.static(m.nv)): + # TODO(team): sparse ten_J + J = ten_J_in[worldid, tenid, i] + if wp.static(m.opt.is_sparse): + efc_J_out[worldid, 0, efcid * wp.static(m.nv) + i] = J + else: efc_J_out[worldid, efcid, i] = J - Jqvel += J * qvel_in[worldid, i] + + Jqvel += J * qvel_in[worldid, i] tendon_invweight0_id = worldid % tendon_invweight0.shape[0] - tendon_solref_lim_id = worldid % tendon_solref_lim.shape[0] - tendon_solimp_lim_id = worldid % tendon_solimp_lim.shape[0] - _update_efc_row( + tendon_solref_fri_id = worldid % tendon_solref_fri.shape[0] + tendon_solimp_fri_id = worldid % tendon_solimp_fri.shape[0] + efc_row( worldid, opt_timestep[worldid % opt_timestep.shape[0]], - refsafe_in, efcid, - pos, - pos, + 0.0, + 0.0, tendon_invweight0[tendon_invweight0_id, tenid], - tendon_solref_lim[tendon_solref_lim_id, tenid], - tendon_solimp_lim[tendon_solimp_lim_id, tenid], - tenmargin, - Jqvel, + tendon_solref_fri[tendon_solref_fri_id, tenid], + tendon_solimp_fri[tendon_solimp_fri_id, tenid], 0.0, - ConstraintType.LIMIT_TENDON, + Jqvel, + frictionloss, + ConstraintType.FRICTION_TENDON, tenid, efc_type_out, efc_id_out, @@ -1206,368 +1144,679 @@ def _efc_limit_tendon( efc_frictionloss_out, ) + @nested_kernel(module="unique", enable_backward=False) + def limit_slide_hinge( + # Model: + opt_timestep: wp.array(dtype=float), + jnt_qposadr: wp.array(dtype=int), + jnt_dofadr: wp.array(dtype=int), + jnt_solref: wp.array2d(dtype=wp.vec2), + jnt_solimp: wp.array2d(dtype=vec5), + jnt_range: wp.array2d(dtype=wp.vec2), + jnt_margin: wp.array2d(dtype=float), + dof_invweight0: wp.array2d(dtype=float), + jnt_limited_slide_hinge_adr: wp.array(dtype=int), + # Data in: + qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), + # Data out: + nl_out: wp.array(dtype=int), + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + worldid, jntlimitedid = wp.tid() + jntid = jnt_limited_slide_hinge_adr[jntlimitedid] + jnt_range_id = worldid % jnt_range.shape[0] + jntrange = jnt_range[jnt_range_id, jntid] + + qpos = qpos_in[worldid, jnt_qposadr[jntid]] + jnt_margin_id = worldid % jnt_margin.shape[0] + jntmargin = jnt_margin[jnt_margin_id, jntid] + dist_min, dist_max = qpos - jntrange[0], jntrange[1] - qpos + pos = wp.min(dist_min, dist_max) - jntmargin + active = pos < 0 + + if active: + wp.atomic_add(nl_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= wp.static(d.njmax): + return + + dofadr = jnt_dofadr[jntid] + + J = float(dist_min < dist_max) * 2.0 - 1.0 + + if wp.static(m.opt.is_sparse): + efc_J_rownnz_out[worldid, efcid] = 1 + rowadr = efcid * wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + efc_J_colind_out[worldid, rowadr] = dofadr + efc_J_out[worldid, 0, rowadr] = J + else: + for i in range(wp.static(m.nv)): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dofadr] = J + + Jqvel = J * qvel_in[worldid, dofadr] -@wp.kernel -def _efc_contact_pyramidal( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - opt_impratio_invsqrt: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - njmax_in: int, - nacon_in: wp.array(dtype=int), - # In: - refsafe_in: int, - dist_in: wp.array(dtype=float), - condim_in: wp.array(dtype=int), - includemargin_in: wp.array(dtype=float), - worldid_in: wp.array(dtype=int), - geom_in: wp.array(dtype=wp.vec2i), - pos_in: wp.array(dtype=wp.vec3), - frame_in: wp.array(dtype=wp.mat33), - friction_in: wp.array(dtype=vec5), - solref_in: wp.array(dtype=wp.vec2), - solimp_in: wp.array(dtype=vec5), - type_in: wp.array(dtype=int), - # Data out: - nefc_out: wp.array(dtype=int), - contact_efc_address_out: wp.array2d(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - conid, dimid = wp.tid() + dof_invweight0_id = worldid % dof_invweight0.shape[0] + jnt_solref_id = worldid % jnt_solref.shape[0] + jnt_solimp_id = worldid % jnt_solimp.shape[0] + efc_row( + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + dof_invweight0[dof_invweight0_id, dofadr], + jnt_solref[jnt_solref_id, jntid], + jnt_solimp[jnt_solimp_id, jntid], + jntmargin, + Jqvel, + 0.0, + ConstraintType.LIMIT_JOINT, + jntid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) - if conid >= nacon_in[0]: - return + @nested_kernel(module="unique", enable_backward=False) + def limit_ball( + # Model: + opt_timestep: wp.array(dtype=float), + jnt_qposadr: wp.array(dtype=int), + jnt_dofadr: wp.array(dtype=int), + jnt_solref: wp.array2d(dtype=wp.vec2), + jnt_solimp: wp.array2d(dtype=vec5), + jnt_range: wp.array2d(dtype=wp.vec2), + jnt_margin: wp.array2d(dtype=float), + dof_invweight0: wp.array2d(dtype=float), + jnt_limited_ball_adr: wp.array(dtype=int), + # Data in: + qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), + # Data out: + nl_out: wp.array(dtype=int), + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + worldid, jntlimitedid = wp.tid() + jntid = jnt_limited_ball_adr[jntlimitedid] + qposadr = jnt_qposadr[jntid] + + qpos = qpos_in[worldid] + jnt_quat = wp.quat(qpos[qposadr + 0], qpos[qposadr + 1], qpos[qposadr + 2], qpos[qposadr + 3]) + jnt_quat = wp.normalize(jnt_quat) + axis_angle = math.quat_to_vel(jnt_quat) + jnt_range_id = worldid % jnt_range.shape[0] + jntrange = jnt_range[jnt_range_id, jntid] + axis, angle = math.normalize_with_norm(axis_angle) + jnt_margin_id = worldid % jnt_margin.shape[0] + jntmargin = jnt_margin[jnt_margin_id, jntid] + + pos = wp.max(jntrange[0], jntrange[1]) - angle - jntmargin + active = pos < 0 + + if active: + wp.atomic_add(nl_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= wp.static(d.njmax): + return + + dofadr = jnt_dofadr[jntid] + dof0 = dofadr + 0 + dof1 = dofadr + 1 + dof2 = dofadr + 2 + + if wp.static(m.opt.is_sparse): + efc_J_rownnz_out[worldid, efcid] = 3 + rowadr = efcid * wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + + sparseid0 = rowadr + 0 + sparseid1 = rowadr + 1 + sparseid2 = rowadr + 2 + + efc_J_colind_out[worldid, sparseid0] = dof0 + efc_J_colind_out[worldid, sparseid1] = dof1 + efc_J_colind_out[worldid, sparseid2] = dof2 + + efc_J_out[worldid, 0, sparseid0] = -axis[0] + efc_J_out[worldid, 0, sparseid1] = -axis[1] + efc_J_out[worldid, 0, sparseid2] = -axis[2] + else: + for i in range(wp.static(m.nv)): + efc_J_out[worldid, efcid, i] = 0.0 + efc_J_out[worldid, efcid, dof0] = -axis[0] + efc_J_out[worldid, efcid, dof1] = -axis[1] + efc_J_out[worldid, efcid, dof2] = -axis[2] + + Jqvel = -axis[0] * qvel_in[worldid, dof0] + Jqvel -= axis[1] * qvel_in[worldid, dof1] + Jqvel -= axis[2] * qvel_in[worldid, dof2] + + dof_invweight0_id = worldid % dof_invweight0.shape[0] + jnt_solref_id = worldid % jnt_solref.shape[0] + jnt_solimp_id = worldid % jnt_solimp.shape[0] + efc_row( + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + dof_invweight0[dof_invweight0_id, dofadr], + jnt_solref[jnt_solref_id, jntid], + jnt_solimp[jnt_solimp_id, jntid], + jntmargin, + Jqvel, + 0.0, + ConstraintType.LIMIT_JOINT, + jntid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) - if not type_in[conid] & ContactType.CONSTRAINT: - return + @nested_kernel(module="unique", enable_backward=False) + def limit_tendon( + # Model: + opt_timestep: wp.array(dtype=float), + jnt_dofadr: wp.array(dtype=int), + tendon_adr: wp.array(dtype=int), + tendon_num: wp.array(dtype=int), + tendon_solref_lim: wp.array2d(dtype=wp.vec2), + tendon_solimp_lim: wp.array2d(dtype=vec5), + tendon_range: wp.array2d(dtype=wp.vec2), + tendon_margin: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), + wrap_type: wp.array(dtype=int), + wrap_objid: wp.array(dtype=int), + tendon_limited_adr: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + ten_J_in: wp.array3d(dtype=float), + ten_length_in: wp.array2d(dtype=float), + # Data out: + nl_out: wp.array(dtype=int), + nefc_out: wp.array(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + worldid, tenlimitedid = wp.tid() + tenid = tendon_limited_adr[tenlimitedid] + + tendon_range_id = worldid % tendon_range.shape[0] + tenrange = tendon_range[tendon_range_id, tenid] + length = ten_length_in[worldid, tenid] + dist_min, dist_max = length - tenrange[0], tenrange[1] - length + tendon_margin_id = worldid % tendon_margin.shape[0] + tenmargin = tendon_margin[tendon_margin_id, tenid] + pos = wp.min(dist_min, dist_max) - tenmargin + active = pos < 0 + + if active: + wp.atomic_add(nl_out, worldid, 1) + efcid = wp.atomic_add(nefc_out, worldid, 1) + + if efcid >= wp.static(d.njmax): + return + + Jqvel = float(0.0) + scl = float(dist_min < dist_max) * 2.0 - 1.0 + + adr = tendon_adr[tenid] + if wrap_type[adr] == types.WrapType.JOINT: + ten_num = tendon_num[tenid] + for i in range(ten_num): + dofadr = jnt_dofadr[wrap_objid[adr + i]] + J = scl * ten_J_in[worldid, tenid, dofadr] + + if wp.static(m.opt.is_sparse): + # TODO(team): efc.J_rownnz, efc.J_rowadr, efc.J_colind + efc_J_out[worldid, 0, efcid * wp.static(m.nv) + dofadr] = J + else: + efc_J_out[worldid, efcid, dofadr] = J - condim = condim_in[conid] + Jqvel += J * qvel_in[worldid, dofadr] + else: + for i in range(wp.static(m.nv)): + J = scl * ten_J_in[worldid, tenid, i] - if condim == 1 and dimid > 0: - return - elif condim > 1 and dimid >= 2 * (condim - 1): - return + if wp.static(m.opt.is_sparse): + # TODO(team): efc.J_rownnz, efc.J_rowadr, efc.J_colind + efc_J_out[worldid, 0, efcid * wp.static(m.nv) + i] = J + else: + efc_J_out[worldid, efcid, i] = J - includemargin = includemargin_in[conid] - pos = dist_in[conid] - includemargin - active = pos < 0 + Jqvel += J * qvel_in[worldid, i] - if active: - worldid = worldid_in[conid] + tendon_invweight0_id = worldid % tendon_invweight0.shape[0] + tendon_solref_lim_id = worldid % tendon_solref_lim.shape[0] + tendon_solimp_lim_id = worldid % tendon_solimp_lim.shape[0] + efc_row( + worldid, + opt_timestep[worldid % opt_timestep.shape[0]], + efcid, + pos, + pos, + tendon_invweight0[tendon_invweight0_id, tenid], + tendon_solref_lim[tendon_solref_lim_id, tenid], + tendon_solimp_lim[tendon_solimp_lim_id, tenid], + tenmargin, + Jqvel, + 0.0, + ConstraintType.LIMIT_TENDON, + tenid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) - efcid = wp.atomic_add(nefc_out, worldid, 1) - if efcid >= njmax_in: - contact_efc_address_out[conid, dimid] = -1 + @nested_kernel(module="unique", enable_backward=False) + def contact_pyramidal( + # Model: + opt_timestep: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), + body_parentid: wp.array(dtype=int), + body_rootid: wp.array(dtype=int), + body_invweight0: wp.array2d(dtype=wp.vec2), + dof_bodyid: wp.array(dtype=int), + geom_bodyid: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cdof_in: wp.array2d(dtype=wp.spatial_vector), + nacon_in: wp.array(dtype=int), + dist_in: wp.array(dtype=float), + condim_in: wp.array(dtype=int), + includemargin_in: wp.array(dtype=float), + worldid_in: wp.array(dtype=int), + geom_in: wp.array(dtype=wp.vec2i), + pos_in: wp.array(dtype=wp.vec3), + frame_in: wp.array(dtype=wp.mat33), + friction_in: wp.array(dtype=vec5), + solref_in: wp.array(dtype=wp.vec2), + solimp_in: wp.array(dtype=vec5), + type_in: wp.array(dtype=int), + # Data out: + nefc_out: wp.array(dtype=int), + contact_efc_address_out: wp.array2d(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + conid, dimid = wp.tid() + + if conid >= nacon_in[0]: return - timestep = opt_timestep[worldid % opt_timestep.shape[0]] - impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] - contact_efc_address_out[conid, dimid] = efcid - - geom = geom_in[conid] - body1 = geom_bodyid[geom[0]] - body2 = geom_bodyid[geom[1]] + if not type_in[conid] & ContactType.CONSTRAINT: + return - con_pos = pos_in[conid] - frame = frame_in[conid] + condim = condim_in[conid] - # pyramidal has common invweight across all edges - body_invweight0_id = worldid % body_invweight0.shape[0] - invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] + if condim == 1 and dimid > 0: + return + elif condim > 1 and dimid >= 2 * (condim - 1): + return - if condim > 1: - dimid2 = dimid / 2 + 1 + includemargin = includemargin_in[conid] + pos = dist_in[conid] - includemargin + active = pos < 0 - friction = friction_in[conid] - fri0 = friction[0] - frii = friction[dimid2 - 1] - invweight = invweight + fri0 * fri0 * invweight - invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt + if active: + worldid = worldid_in[conid] - Jqvel = float(0.0) - for i in range(nv): - J = float(0.0) - Ji = float(0.0) - jac1p, jac1r = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - con_pos, - body1, - i, - worldid, - ) - jac2p, jac2r = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - con_pos, - body2, - i, - worldid, - ) - jacp_dif = jac2p - jac1p - for xyz in range(3): - J += frame[0, xyz] * jacp_dif[xyz] + efcid = wp.atomic_add(nefc_out, worldid, 1) + if efcid >= wp.static(d.njmax): + contact_efc_address_out[conid, dimid] = -1 + return - if condim > 1: - if dimid2 < 3: - Ji += frame[dimid2, xyz] * jacp_dif[xyz] - else: - Ji += frame[dimid2 - 3, xyz] * (jac2r[xyz] - jac1r[xyz]) + timestep = opt_timestep[worldid % opt_timestep.shape[0]] + impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] + contact_efc_address_out[conid, dimid] = efcid - if condim > 1: - if dimid % 2 == 0: - J += Ji * frii - else: - J -= Ji * frii + geom = geom_in[conid] + body1 = geom_bodyid[geom[0]] + body2 = geom_bodyid[geom[1]] - efc_J_out[worldid, efcid, i] = J - Jqvel += J * qvel_in[worldid, i] + con_pos = pos_in[conid] + frame = frame_in[conid] - if condim == 1: - efc_type = ConstraintType.CONTACT_FRICTIONLESS - else: - efc_type = ConstraintType.CONTACT_PYRAMIDAL + # pyramidal has common invweight across all edges + body_invweight0_id = worldid % body_invweight0.shape[0] + invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcid, - pos, - pos, - invweight, - solref_in[conid], - solimp_in[conid], - includemargin, - Jqvel, - 0.0, - efc_type, - conid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + if condim > 1: + dimid2 = dimid / 2 + 1 + friction = friction_in[conid] + fri0 = friction[0] + frii = friction[dimid2 - 1] + invweight = invweight + fri0 * fri0 * invweight + invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt + + if wp.static(m.opt.is_sparse): + efc_J_rownnz_out[worldid, efcid] = wp.static(m.nv) + rowadr = efcid * wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + + Jqvel = float(0.0) + for i in range(wp.static(m.nv)): + J = float(0.0) + Ji = float(0.0) + jac1p, jac1r = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + con_pos, + body1, + i, + worldid, + ) + jac2p, jac2r = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + con_pos, + body2, + i, + worldid, + ) + jacp_dif = jac2p - jac1p + for xyz in range(3): + J += frame[0, xyz] * jacp_dif[xyz] -@wp.kernel -def _efc_contact_elliptic( - # Model: - nv: int, - opt_timestep: wp.array(dtype=float), - opt_impratio_invsqrt: wp.array(dtype=float), - body_parentid: wp.array(dtype=int), - body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=wp.vec2), - dof_bodyid: wp.array(dtype=int), - geom_bodyid: wp.array(dtype=int), - # Data in: - qvel_in: wp.array2d(dtype=float), - subtree_com_in: wp.array2d(dtype=wp.vec3), - cdof_in: wp.array2d(dtype=wp.spatial_vector), - njmax_in: int, - nacon_in: wp.array(dtype=int), - # In: - refsafe_in: int, - dist_in: wp.array(dtype=float), - condim_in: wp.array(dtype=int), - includemargin_in: wp.array(dtype=float), - worldid_in: wp.array(dtype=int), - geom_in: wp.array(dtype=wp.vec2i), - pos_in: wp.array(dtype=wp.vec3), - frame_in: wp.array(dtype=wp.mat33), - friction_in: wp.array(dtype=vec5), - solref_in: wp.array(dtype=wp.vec2), - solreffriction_in: wp.array(dtype=wp.vec2), - solimp_in: wp.array(dtype=vec5), - type_in: wp.array(dtype=int), - # Data out: - nefc_out: wp.array(dtype=int), - contact_efc_address_out: wp.array2d(dtype=int), - efc_type_out: wp.array2d(dtype=int), - efc_id_out: wp.array2d(dtype=int), - efc_J_out: wp.array3d(dtype=float), - efc_pos_out: wp.array2d(dtype=float), - efc_margin_out: wp.array2d(dtype=float), - efc_D_out: wp.array2d(dtype=float), - efc_vel_out: wp.array2d(dtype=float), - efc_aref_out: wp.array2d(dtype=float), - efc_frictionloss_out: wp.array2d(dtype=float), -): - conid, dimid = wp.tid() + if condim > 1: + if dimid2 < 3: + Ji += frame[dimid2, xyz] * jacp_dif[xyz] + else: + Ji += frame[dimid2 - 3, xyz] * (jac2r[xyz] - jac1r[xyz]) - if conid >= nacon_in[0]: - return + if condim > 1: + if dimid % 2 == 0: + J += Ji * frii + else: + J -= Ji * frii - if not type_in[conid] & ContactType.CONSTRAINT: - return + if wp.static(m.opt.is_sparse): + sparseid = rowadr + i + efc_J_colind_out[worldid, sparseid] = i + efc_J_out[worldid, 0, sparseid] = J + else: + efc_J_out[worldid, efcid, i] = J - condim = condim_in[conid] + Jqvel += J * qvel_in[worldid, i] - if dimid > condim - 1: - return + if condim == 1: + efc_type = ConstraintType.CONTACT_FRICTIONLESS + else: + efc_type = ConstraintType.CONTACT_PYRAMIDAL - includemargin = includemargin_in[conid] - pos = dist_in[conid] - includemargin - active = pos < 0.0 + efc_row( + worldid, + timestep, + efcid, + pos, + pos, + invweight, + solref_in[conid], + solimp_in[conid], + includemargin, + Jqvel, + 0.0, + efc_type, + conid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) - if active: - worldid = worldid_in[conid] + @nested_kernel(module="unique", enable_backward=False) + def contact_elliptic( + # Model: + opt_timestep: wp.array(dtype=float), + opt_impratio_invsqrt: wp.array(dtype=float), + body_parentid: wp.array(dtype=int), + body_rootid: wp.array(dtype=int), + body_invweight0: wp.array2d(dtype=wp.vec2), + dof_bodyid: wp.array(dtype=int), + geom_bodyid: wp.array(dtype=int), + # Data in: + qvel_in: wp.array2d(dtype=float), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cdof_in: wp.array2d(dtype=wp.spatial_vector), + nacon_in: wp.array(dtype=int), + dist_in: wp.array(dtype=float), + condim_in: wp.array(dtype=int), + includemargin_in: wp.array(dtype=float), + worldid_in: wp.array(dtype=int), + geom_in: wp.array(dtype=wp.vec2i), + pos_in: wp.array(dtype=wp.vec3), + frame_in: wp.array(dtype=wp.mat33), + friction_in: wp.array(dtype=vec5), + solref_in: wp.array(dtype=wp.vec2), + solreffriction_in: wp.array(dtype=wp.vec2), + solimp_in: wp.array(dtype=vec5), + type_in: wp.array(dtype=int), + # Data out: + nefc_out: wp.array(dtype=int), + contact_efc_address_out: wp.array2d(dtype=int), + efc_type_out: wp.array2d(dtype=int), + efc_id_out: wp.array2d(dtype=int), + efc_J_rownnz_out: wp.array2d(dtype=int), + efc_J_rowadr_out: wp.array2d(dtype=int), + efc_J_colind_out: wp.array2d(dtype=int), + efc_J_out: wp.array3d(dtype=float), + efc_pos_out: wp.array2d(dtype=float), + efc_margin_out: wp.array2d(dtype=float), + efc_D_out: wp.array2d(dtype=float), + efc_vel_out: wp.array2d(dtype=float), + efc_aref_out: wp.array2d(dtype=float), + efc_frictionloss_out: wp.array2d(dtype=float), + ): + conid, dimid = wp.tid() + + if conid >= nacon_in[0]: + return - efcid = wp.atomic_add(nefc_out, worldid, 1) - if efcid >= njmax_in: - contact_efc_address_out[conid, dimid] = -1 + if not type_in[conid] & ContactType.CONSTRAINT: return - timestep = opt_timestep[worldid % opt_timestep.shape[0]] - impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] - contact_efc_address_out[conid, dimid] = efcid + condim = condim_in[conid] - geom = geom_in[conid] - body1 = geom_bodyid[geom[0]] - body2 = geom_bodyid[geom[1]] + if dimid > condim - 1: + return - cpos = pos_in[conid] - frame = frame_in[conid] + includemargin = includemargin_in[conid] + pos = dist_in[conid] - includemargin + active = pos < 0.0 + + if active: + worldid = worldid_in[conid] + + efcid = wp.atomic_add(nefc_out, worldid, 1) + if efcid >= wp.static(d.njmax): + contact_efc_address_out[conid, dimid] = -1 + return + + timestep = opt_timestep[worldid % opt_timestep.shape[0]] + impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] + contact_efc_address_out[conid, dimid] = efcid + + geom = geom_in[conid] + body1 = geom_bodyid[geom[0]] + body2 = geom_bodyid[geom[1]] + + cpos = pos_in[conid] + frame = frame_in[conid] + + if wp.static(m.opt.is_sparse): + efc_J_rownnz_out[worldid, efcid] = wp.static(m.nv) + rowadr = efcid * wp.static(m.nv) + efc_J_rowadr_out[worldid, efcid] = rowadr + + # TODO(team): parallelize J and Jqvel computation? + Jqvel = float(0.0) + for i in range(wp.static(m.nv)): + J = float(0.0) + jac1p, jac1r = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + cpos, + body1, + i, + worldid, + ) + jac2p, jac2r = support.jac( + body_parentid, + body_rootid, + dof_bodyid, + subtree_com_in, + cdof_in, + cpos, + body2, + i, + worldid, + ) + for xyz in range(3): + if dimid < 3: + jac_dif = jac2p[xyz] - jac1p[xyz] + J += frame[dimid, xyz] * jac_dif + else: + jac_dif = jac2r[xyz] - jac1r[xyz] + J += frame[dimid - 3, xyz] * jac_dif - # TODO(team): parallelize J and Jqvel computation? - Jqvel = float(0.0) - for i in range(nv): - J = float(0.0) - jac1p, jac1r = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - cpos, - body1, - i, - worldid, - ) - jac2p, jac2r = support.jac( - body_parentid, - body_rootid, - dof_bodyid, - subtree_com_in, - cdof_in, - cpos, - body2, - i, - worldid, - ) - for xyz in range(3): - if dimid < 3: - jac_dif = jac2p[xyz] - jac1p[xyz] - J += frame[dimid, xyz] * jac_dif + if wp.static(m.opt.is_sparse): + sparseid = rowadr + i + efc_J_colind_out[worldid, sparseid] = i + efc_J_out[worldid, 0, sparseid] = J else: - jac_dif = jac2r[xyz] - jac1r[xyz] - J += frame[dimid - 3, xyz] * jac_dif - - efc_J_out[worldid, efcid, i] = J - Jqvel += J * qvel_in[worldid, i] - - body_invweight0_id = worldid % body_invweight0.shape[0] - invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] + efc_J_out[worldid, efcid, i] = J - ref = solref_in[conid] - pos_aref = pos + Jqvel += J * qvel_in[worldid, i] - if dimid > 0: - solreffriction = solreffriction_in[conid] + body_invweight0_id = worldid % body_invweight0.shape[0] + invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0] - # non-normal directions use solreffriction (if non-zero) - if solreffriction[0] or solreffriction[1]: - ref = solreffriction + ref = solref_in[conid] + pos_aref = pos - invweight = invweight * impratio_invsqrt * impratio_invsqrt - friction = friction_in[conid] + if dimid > 0: + solreffriction = solreffriction_in[conid] - if dimid > 1: - fri0 = friction[0] - frii = friction[dimid - 1] - fri = fri0 * fri0 / (frii * frii) - invweight *= fri + # non-normal directions use solreffriction (if non-zero) + if solreffriction[0] or solreffriction[1]: + ref = solreffriction - pos_aref = 0.0 + invweight = invweight * impratio_invsqrt * impratio_invsqrt + friction = friction_in[conid] - if condim == 1: - efc_type = ConstraintType.CONTACT_FRICTIONLESS - else: - efc_type = ConstraintType.CONTACT_ELLIPTIC - - _update_efc_row( - worldid, - timestep, - refsafe_in, - efcid, - pos_aref, - pos, - invweight, - ref, - solimp_in[conid], - includemargin, - Jqvel, - 0.0, - efc_type, - conid, - efc_type_out, - efc_id_out, - efc_pos_out, - efc_margin_out, - efc_D_out, - efc_vel_out, - efc_aref_out, - efc_frictionloss_out, - ) + if dimid > 1: + fri0 = friction[0] + frii = friction[dimid - 1] + fri = fri0 * fri0 / (frii * frii) + invweight *= fri + pos_aref = 0.0 -@wp.kernel -def _num_equality( - # Data in: - ne_connect_in: wp.array(dtype=int), - ne_weld_in: wp.array(dtype=int), - ne_jnt_in: wp.array(dtype=int), - ne_ten_in: wp.array(dtype=int), - ne_flex_in: wp.array(dtype=int), - # Data out: - ne_out: wp.array(dtype=int), -): - worldid = wp.tid() - ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne_flex_in[worldid] - ne_out[worldid] = ne + if condim == 1: + efc_type = ConstraintType.CONTACT_FRICTIONLESS + else: + efc_type = ConstraintType.CONTACT_ELLIPTIC + efc_row( + worldid, + timestep, + efcid, + pos_aref, + pos, + invweight, + ref, + solimp_in[conid], + includemargin, + Jqvel, + 0.0, + efc_type, + conid, + efc_type_out, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_vel_out, + efc_aref_out, + efc_frictionloss_out, + ) -@event_scope -def make_constraint(m: types.Model, d: types.Data): - """Creates constraint jacobians and other supporting data.""" wp.launch( _zero_constraint_counts, dim=d.nworld, @@ -1575,20 +1824,27 @@ def make_constraint(m: types.Model, d: types.Data): ) if not (m.opt.disableflags & types.DisableBit.CONSTRAINT): - refsafe = m.opt.disableflags & types.DisableBit.REFSAFE + if m.opt.is_sparse: + wp.launch( + efc_J_initialize, + dim=(d.nworld, d.njmax, m.nv), + inputs=[], + outputs=[d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J], + ) if not (m.opt.disableflags & types.DisableBit.EQUALITY): wp.launch( - _efc_equality_connect, + equality_connect, dim=(d.nworld, m.eq_connect_adr.size), inputs=[ - m.nv, - m.nsite, m.opt.timestep, m.body_parentid, m.body_rootid, + m.body_dofnum, + m.body_dofadr, m.body_invweight0, m.dof_bodyid, + m.dof_parentid, m.site_bodyid, m.eq_obj1id, m.eq_obj2id, @@ -1604,13 +1860,14 @@ def make_constraint(m: types.Model, d: types.Data): d.site_xpos, d.subtree_com, d.cdof, - d.njmax, - refsafe, ], outputs=[ d.nefc, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1622,16 +1879,17 @@ def make_constraint(m: types.Model, d: types.Data): ], ) wp.launch( - _efc_equality_weld, + equality_weld, dim=(d.nworld, m.eq_wld_adr.size), inputs=[ - m.nv, - m.nsite, m.opt.timestep, m.body_parentid, m.body_rootid, + m.body_dofnum, + m.body_dofadr, m.body_invweight0, m.dof_bodyid, + m.dof_parentid, m.site_bodyid, m.site_quat, m.eq_obj1id, @@ -1649,13 +1907,14 @@ def make_constraint(m: types.Model, d: types.Data): d.site_xpos, d.subtree_com, d.cdof, - d.njmax, - refsafe, ], outputs=[ d.nefc, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1667,10 +1926,9 @@ def make_constraint(m: types.Model, d: types.Data): ], ) wp.launch( - _efc_equality_joint, + equality_joint, dim=(d.nworld, m.eq_jnt_adr.size), inputs=[ - m.nv, m.opt.timestep, m.qpos0, m.jnt_qposadr, @@ -1685,13 +1943,14 @@ def make_constraint(m: types.Model, d: types.Data): d.qpos, d.qvel, d.eq_active, - d.njmax, - refsafe, ], outputs=[ d.nefc, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1703,10 +1962,9 @@ def make_constraint(m: types.Model, d: types.Data): ], ) wp.launch( - _efc_equality_tendon, + equality_tendon, dim=(d.nworld, m.eq_ten_adr.size), inputs=[ - m.nv, m.opt.timestep, m.eq_obj1id, m.eq_obj2id, @@ -1720,8 +1978,6 @@ def make_constraint(m: types.Model, d: types.Data): d.eq_active, d.ten_J, d.ten_length, - d.njmax, - refsafe, ], outputs=[ d.nefc, @@ -1739,10 +1995,9 @@ def make_constraint(m: types.Model, d: types.Data): ) wp.launch( - _efc_equality_flex, + equality_flex, dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge), inputs=[ - m.nv, m.opt.timestep, m.flexedge_length0, m.flexedge_invweight0, @@ -1752,8 +2007,6 @@ def make_constraint(m: types.Model, d: types.Data): d.qvel, d.flexedge_J, d.flexedge_length, - d.njmax, - refsafe, ], outputs=[ d.nefc, @@ -1779,24 +2032,24 @@ def make_constraint(m: types.Model, d: types.Data): if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS): wp.launch( - _efc_friction_dof, + friction_dof, dim=(d.nworld, m.nv), inputs=[ - m.nv, m.opt.timestep, m.dof_solref, m.dof_solimp, m.dof_frictionloss, m.dof_invweight0, d.qvel, - d.njmax, - refsafe, ], outputs=[ d.nf, d.nefc, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1808,10 +2061,9 @@ def make_constraint(m: types.Model, d: types.Data): ) wp.launch( - _efc_friction_tendon, + friction_tendon, dim=(d.nworld, m.ntendon), inputs=[ - m.nv, m.opt.timestep, m.tendon_solref_fri, m.tendon_solimp_fri, @@ -1819,8 +2071,6 @@ def make_constraint(m: types.Model, d: types.Data): m.tendon_invweight0, d.qvel, d.ten_J, - d.njmax, - refsafe, ], outputs=[ d.nf, @@ -1840,10 +2090,9 @@ def make_constraint(m: types.Model, d: types.Data): # limit if not (m.opt.disableflags & types.DisableBit.LIMIT): wp.launch( - _efc_limit_ball, + limit_ball, dim=(d.nworld, m.jnt_limited_ball_adr.size), inputs=[ - m.nv, m.opt.timestep, m.jnt_qposadr, m.jnt_dofadr, @@ -1855,14 +2104,15 @@ def make_constraint(m: types.Model, d: types.Data): m.jnt_limited_ball_adr, d.qpos, d.qvel, - d.njmax, - refsafe, ], outputs=[ d.nl, d.nefc, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1874,10 +2124,9 @@ def make_constraint(m: types.Model, d: types.Data): ) wp.launch( - _efc_limit_slide_hinge, + limit_slide_hinge, dim=(d.nworld, m.jnt_limited_slide_hinge_adr.size), inputs=[ - m.nv, m.opt.timestep, m.jnt_qposadr, m.jnt_dofadr, @@ -1889,14 +2138,15 @@ def make_constraint(m: types.Model, d: types.Data): m.jnt_limited_slide_hinge_adr, d.qpos, d.qvel, - d.njmax, - refsafe, ], outputs=[ d.nl, d.nefc, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1908,10 +2158,9 @@ def make_constraint(m: types.Model, d: types.Data): ) wp.launch( - _efc_limit_tendon, + limit_tendon, dim=(d.nworld, m.tendon_limited_adr.size), inputs=[ - m.nv, m.opt.timestep, m.jnt_dofadr, m.tendon_adr, @@ -1927,8 +2176,6 @@ def make_constraint(m: types.Model, d: types.Data): d.qvel, d.ten_J, d.ten_length, - d.njmax, - refsafe, ], outputs=[ d.nl, @@ -1949,10 +2196,9 @@ def make_constraint(m: types.Model, d: types.Data): if not (m.opt.disableflags & types.DisableBit.CONTACT): if m.opt.cone == types.ConeType.PYRAMIDAL: wp.launch( - _efc_contact_pyramidal, + contact_pyramidal, dim=(d.naconmax, m.nmaxpyramid), inputs=[ - m.nv, m.opt.timestep, m.opt.impratio_invsqrt, m.body_parentid, @@ -1963,9 +2209,7 @@ def make_constraint(m: types.Model, d: types.Data): d.qvel, d.subtree_com, d.cdof, - d.njmax, d.nacon, - refsafe, d.contact.dist, d.contact.dim, d.contact.includemargin, @@ -1983,6 +2227,9 @@ def make_constraint(m: types.Model, d: types.Data): d.contact.efc_address, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, @@ -1994,10 +2241,9 @@ def make_constraint(m: types.Model, d: types.Data): ) elif m.opt.cone == types.ConeType.ELLIPTIC: wp.launch( - _efc_contact_elliptic, + contact_elliptic, dim=(d.naconmax, m.nmaxcondim), inputs=[ - m.nv, m.opt.timestep, m.opt.impratio_invsqrt, m.body_parentid, @@ -2008,9 +2254,7 @@ def make_constraint(m: types.Model, d: types.Data): d.qvel, d.subtree_com, d.cdof, - d.njmax, d.nacon, - refsafe, d.contact.dist, d.contact.dim, d.contact.includemargin, @@ -2029,6 +2273,9 @@ def make_constraint(m: types.Model, d: types.Data): d.contact.efc_address, d.efc.type, d.efc.id, + d.efc.J_rownnz, + d.efc.J_rowadr, + d.efc.J_colind, d.efc.J, d.efc.pos, d.efc.margin, diff --git a/mujoco_warp/_src/constraint_test.py b/mujoco_warp/_src/constraint_test.py index 87c4a4871..6094ad0a5 100644 --- a/mujoco_warp/_src/constraint_test.py +++ b/mujoco_warp/_src/constraint_test.py @@ -36,7 +36,7 @@ def _assert_eq(a, b, name): np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) -def _assert_efc_eq(d, mjd, nefc, name, nv): +def _assert_efc_eq(mjm, m, d, mjd, nefc, name, nv): """Assert equality of efc fields after sorting both sides.""" # Get the ordering indices based on efc_type, efc_pos, efc_vel, efc_aref, efc_d for MJWarp efc_type = d.efc.type.numpy()[0, :nefc] @@ -55,17 +55,36 @@ def _assert_efc_eq(d, mjd, nefc, name, nv): d_sort_indices = np.lexsort((efc_pos, efc_type, efc_vel, efc_aref, efc_d)) mjd_sort_indices = np.lexsort((mjd_efc_pos, mjd_efc_type, mjd_efc_vel, mjd_efc_aref, mjd_efc_d)) + # convert sparse to dense if necessary + if m.opt.is_sparse: + efc_J = np.zeros((nefc, nv)) + mujoco.mju_sparse2dense( + efc_J, + d.efc.J.numpy()[0, 0], + d.efc.J_rownnz.numpy()[0, :nefc], + d.efc.J_rowadr.numpy()[0, :nefc], + d.efc.J_colind.numpy()[0], + ) + else: + efc_J = d.efc.J.numpy()[0] + # Sort MJWarp efc fields - d_sorted = d.efc.J.numpy()[0, d_sort_indices, :nv].reshape(-1) + d_sorted = efc_J[d_sort_indices, :nv].reshape(-1) # Sort MuJoCo efc fields # For J matrix, need to reshape to 2D, sort rows, then flatten nefc = len(mjd_sort_indices) + + if mujoco.mj_isSparse(mjm): + mj_efc_J = np.zeros((mjd.nefc, mjm.nv)) + mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind) + else: + mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv)) + if nv > 0: - mjd_J_2d = mjd.efc_J.reshape(nefc, nv) - mjd_sorted_J = mjd_J_2d[mjd_sort_indices].reshape(-1) + mjd_sorted_J = mj_efc_J[mjd_sort_indices].reshape(-1) else: - mjd_sorted_J = mjd.efc_J + mjd_sorted_J = mj_efc_J mjd_sorted_D = mjd.efc_D[mjd_sort_indices] mjd_sorted_vel = mjd.efc_vel[mjd_sort_indices] @@ -163,7 +182,7 @@ def test_condim(self, cone, condim1, condim2): def test_constraints(self, cone): """Test constraints.""" for key in range(3): - _, mjd, m, d = test_data.fixture("constraints.xml", keyframe=key, overrides={"opt.cone": cone}) + mjm, mjd, m, d = test_data.fixture("constraints.xml", keyframe=key, overrides={"opt.cone": cone}) for arr in (d.ne, d.nefc, d.nf, d.nl, d.efc.type): arr.fill_(-1) @@ -176,12 +195,12 @@ def test_constraints(self, cone): _assert_eq(d.nefc.numpy()[0], mjd.nefc, "nefc") _assert_eq(d.nf.numpy()[0], mjd.nf, "nf") _assert_eq(d.nl.numpy()[0], mjd.nl, "nl") - _assert_efc_eq(d, mjd, mjd.nefc, "efc", m.nv) + _assert_efc_eq(mjm, m, d, mjd, mjd.nefc, "efc", m.nv) def test_limit_tendon(self): """Test limit tendon constraints.""" for keyframe in range(-1, 1): - _, mjd, m, d = test_data.fixture("tendon/tendon_limit.xml", keyframe=keyframe) + mjm, mjd, m, d = test_data.fixture("tendon/tendon_limit.xml", keyframe=keyframe) for arr in (d.nefc, d.nl, d.efc.type): arr.fill_(-1) @@ -192,11 +211,11 @@ def test_limit_tendon(self): _assert_eq(d.nefc.numpy()[0], mjd.nefc, "nefc") _assert_eq(d.nl.numpy()[0], mjd.nl, "nl") - _assert_efc_eq(d, mjd, mjd.nefc, "efc", m.nv) + _assert_efc_eq(mjm, m, d, mjd, mjd.nefc, "efc", m.nv) def test_equality_tendon(self): """Test equality tendon constraints.""" - _, mjd, m, d = test_data.fixture( + mjm, mjd, m, d = test_data.fixture( xml="""