diff --git a/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py b/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py index dc2326dd..42d0440c 100644 --- a/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py +++ b/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py @@ -136,8 +136,8 @@ def test_all_issues( @kernel def test_no_issues( # Model: - qpos0: wp.array(dtype=wp.float32, ndim=1), - geom_pos: wp.array(dtype=wp.vec3, ndim=1), + qpos0: wp.array(dtype=wp.float32, ndim=2), + geom_pos: wp.array(dtype=wp.vec3, ndim=2), # Data in: qpos_in: wp.array(dtype=wp.float32, ndim=2), qvel_in: wp.array(dtype=wp.float32, ndim=2), diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index ae3f2849..1ab37e74 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -210,7 +210,7 @@ def box_box_kernel( worldid = d.collision_worldid[bp_idx] geoms, margin, gap, condim, friction, solref, solreffriction, solimp = ( - contact_params(m, d, tid) + contact_params(m, d, tid, worldid) ) # transformations @@ -220,8 +220,8 @@ def box_box_kernel( trans_atob = b_mat_inv @ (a_pos - b_pos) rot_atob = b_mat_inv @ a_mat - a_size = m.geom_size[ga] - b_size = m.geom_size[gb] + a_size = m.geom_size[worldid, ga] + b_size = m.geom_size[worldid, gb] a = box(rot_atob, trans_atob, a_size) b = box(wp.identity(3, wp.float32), wp.vec3(0.0), b_size) @@ -312,7 +312,10 @@ def box_box_kernel( for i in range(4): pos[i] = pos[idx] - margin = wp.max(m.geom_margin[ga], m.geom_margin[gb]) + margin = wp.max( + m.geom_margin[worldid, ga], + m.geom_margin[worldid, gb], + ) for i in range(4): pos_glob = b_mat @ pos[i] + b_pos n_glob = b_mat @ sep_axis diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index befb2c1c..87a1eda9 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -754,7 +754,7 @@ def _gjk_epa_sparse(m: Model, d: Data): worldid = d.collision_worldid[tid] geoms, margin, gap, condim, friction, solref, solreffriction, solimp = ( - contact_params(m, d, tid) + contact_params(m, d, tid, worldid) ) g1 = geoms[0] @@ -763,10 +763,13 @@ def _gjk_epa_sparse(m: Model, d: Data): if m.geom_type[g1] != geomtype1 or m.geom_type[g2] != geomtype2: return - geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid]) - geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid]) + geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid) + geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid) - margin = wp.max(m.geom_margin[g1], m.geom_margin[g2]) + margin = wp.max( + m.geom_margin[worldid, g1], + m.geom_margin[worldid, g2], + ) simplex, normal = _gjk(m, geom1, geom2) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index 00fdc5b1..db0d72f3 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -32,12 +32,12 @@ @wp.func def _sphere_filter(m: Model, d: Data, geom1: int, geom2: int, worldid: int) -> bool: - margin1 = m.geom_margin[geom1] - margin2 = m.geom_margin[geom2] + margin1 = m.geom_margin[worldid, geom1] + margin2 = m.geom_margin[worldid, geom2] pos1 = d.geom_xpos[worldid, geom1] pos2 = d.geom_xpos[worldid, geom2] - size1 = m.geom_rbound[geom1] - size2 = m.geom_rbound[geom2] + size1 = m.geom_rbound[worldid, geom1] + size2 = m.geom_rbound[worldid, geom2] bound = size1 + size2 + wp.max(margin1, margin2) dif = pos2 - pos1 @@ -105,13 +105,13 @@ def _sap_project(m: Model, d: Data, direction: wp.vec3): worldid, geomid = wp.tid() xpos = d.geom_xpos[worldid, geomid] - rbound = m.geom_rbound[geomid] + rbound = m.geom_rbound[worldid, geomid] if rbound == 0.0: # geom is a plane rbound = MJ_MAXVAL - radius = rbound + m.geom_margin[geomid] + radius = rbound + m.geom_margin[worldid, geomid] center = wp.dot(direction, xpos) d.sap_projection_lower[worldid, geomid] = center - radius diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index ffef2655..e5715a5a 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -44,14 +44,15 @@ def _geom( m: Model, geom_xpos: wp.array(dtype=wp.vec3), geom_xmat: wp.array(dtype=wp.mat33), + worldid: int, ) -> Geom: geom = Geom() geom.pos = geom_xpos[gid] rot = geom_xmat[gid] geom.rot = rot - geom.size = m.geom_size[gid] + geom.size = m.geom_size[worldid, gid] geom.normal = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2]) # plane - dataid = m.geom_dataid[gid] + dataid = m.geom_dataid[worldid, gid] if dataid >= 0: geom.vertadr = m.mesh_vertadr[dataid] geom.vertnum = m.mesh_vertnum[dataid] @@ -737,27 +738,27 @@ def plane_cylinder( @wp.func -def contact_params(m: Model, d: Data, cid: int): +def contact_params(m: Model, d: Data, cid: int, worldid: int): geoms = d.collision_pair[cid] pairid = d.collision_pairid[cid] if pairid > -1: - margin = m.pair_margin[pairid] - gap = m.pair_gap[pairid] + margin = m.pair_margin[worldid, pairid] + gap = m.pair_gap[worldid, pairid] condim = m.pair_dim[pairid] - friction = m.pair_friction[pairid] - solref = m.pair_solref[pairid] - solreffriction = m.pair_solreffriction[pairid] - solimp = m.pair_solimp[pairid] + friction = m.pair_friction[worldid, pairid] + solref = m.pair_solref[worldid, pairid] + solreffriction = m.pair_solreffriction[worldid, pairid] + solimp = m.pair_solimp[worldid, pairid] else: g1 = geoms[0] g2 = geoms[1] - p1 = m.geom_priority[g1] - p2 = m.geom_priority[g2] + p1 = m.geom_priority[worldid, g1] + p2 = m.geom_priority[worldid, g2] - solmix1 = m.geom_solmix[g1] - solmix2 = m.geom_solmix[g2] + solmix1 = m.geom_solmix[worldid, g1] + solmix2 = m.geom_solmix[worldid, g2] mix = solmix1 / (solmix1 + solmix2) mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix) @@ -765,8 +766,14 @@ def contact_params(m: Model, d: Data, cid: int): mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix) mix = wp.where(p1 == p2, mix, wp.where(p1 > p2, 1.0, 0.0)) - margin = wp.max(m.geom_margin[g1], m.geom_margin[g2]) - gap = wp.max(m.geom_gap[g1], m.geom_gap[g2]) + margin = wp.max( + m.geom_margin[worldid, g1], + m.geom_margin[worldid, g2], + ) + gap = wp.max( + m.geom_gap[worldid, g1], + m.geom_gap[worldid, g2], + ) condim1 = m.geom_condim[g1] condim2 = m.geom_condim[g2] @@ -774,7 +781,10 @@ def contact_params(m: Model, d: Data, cid: int): p1 == p2, wp.max(condim1, condim2), wp.where(p1 > p2, condim1, condim2) ) - geom_friction = wp.max(m.geom_friction[g1], m.geom_friction[g2]) + geom_friction = wp.max( + m.geom_friction[worldid, g1], + m.geom_friction[worldid, g2], + ) friction = vec5( geom_friction[0], geom_friction[0], @@ -783,14 +793,19 @@ def contact_params(m: Model, d: Data, cid: int): geom_friction[2], ) - if m.geom_solref[g1].x > 0.0 and m.geom_solref[g2].x > 0.0: - solref = mix * m.geom_solref[g1] + (1.0 - mix) * m.geom_solref[g2] + if m.geom_solref[worldid, g1].x > 0.0 and m.geom_solref[worldid, g2].x > 0.0: + solref = ( + mix * m.geom_solref[worldid, g1] + (1.0 - mix) * m.geom_solref[worldid, g2] + ) else: - solref = wp.min(m.geom_solref[g1], m.geom_solref[g2]) + solref = wp.min( + m.geom_solref[worldid, g1], + m.geom_solref[worldid, g2], + ) solreffriction = wp.vec2(0.0, 0.0) - solimp = mix * m.geom_solimp[g1] + (1.0 - mix) * m.geom_solimp[g2] + solimp = mix * m.geom_solimp[worldid, g1] + (1.0 - mix) * m.geom_solimp[worldid, g2] return geoms, margin, gap, condim, friction, solref, solreffriction, solimp @@ -1258,16 +1273,16 @@ def _primitive_narrowphase( if tid >= d.ncollision[0]: return + worldid = d.collision_worldid[tid] + geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( - m, d, tid + m, d, tid, worldid ) g1 = geoms[0] g2 = geoms[1] - worldid = d.collision_worldid[tid] - - geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid]) - geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid]) + geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid) + geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid], worldid) type1 = m.geom_type[g1] type2 = m.geom_type[g2] diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index efaa02f9..436c2c1b 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -98,7 +98,7 @@ def _efc_equality_connect( for i in range(wp.static(3)): d.efc.worldid[efcid + i] = worldid - data = m.eq_data[i_eq] + data = m.eq_data[worldid, i_eq] anchor1 = wp.vec3f(data[0], data[1], data[2]) anchor2 = wp.vec3f(data[3], data[4], data[5]) @@ -131,11 +131,13 @@ def _efc_equality_connect( d.efc.J[efcid + 2, dofid] = j1mj2[2] Jqvel += j1mj2 * d.qvel[worldid, dofid] - invweight = m.body_invweight0[body1id, 0] + m.body_invweight0[body2id, 0] + invweight = ( + m.body_invweight0[worldid, body1id, 0] + m.body_invweight0[worldid, body2id, 0] + ) pos_imp = wp.length(pos) - solref = m.eq_solref[i_eq] - solimp = m.eq_solimp[i_eq] + solref = m.eq_solref[worldid, i_eq] + solimp = m.eq_solimp[worldid, i_eq] for i in range(3): efcidi = efcid + i @@ -173,7 +175,7 @@ def _efc_equality_joint( jntid_1 = m.eq_obj1id[i_eq] jntid_2 = m.eq_obj2id[i_eq] - data = m.eq_data[i_eq] + data = m.eq_data[worldid, i_eq] dofadr1 = m.jnt_dofadr[jntid_1] qposadr1 = m.jnt_qposadr[jntid_1] d.efc.J[efcid, dofadr1] = 1.0 @@ -182,7 +184,7 @@ def _efc_equality_joint( # Two joint constraint qposadr2 = m.jnt_qposadr[jntid_2] dofadr2 = m.jnt_dofadr[jntid_2] - dif = d.qpos[worldid, qposadr2] - m.qpos0[qposadr2] + dif = d.qpos[worldid, qposadr2] - m.qpos0[worldid, qposadr2] # Horner's method for polynomials rhs = data[0] + dif * (data[1] + dif * (data[2] + dif * (data[3] + dif * data[4]))) @@ -190,16 +192,16 @@ def _efc_equality_joint( 2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4]) ) - pos = d.qpos[worldid, qposadr1] - m.qpos0[qposadr1] - rhs + pos = d.qpos[worldid, qposadr1] - m.qpos0[worldid, qposadr1] - rhs Jqvel = d.qvel[worldid, dofadr1] - d.qvel[worldid, dofadr2] * deriv_2 - invweight = m.dof_invweight0[dofadr1] + m.dof_invweight0[dofadr2] + invweight = m.dof_invweight0[worldid, dofadr1] + m.dof_invweight0[worldid, dofadr2] d.efc.J[efcid, dofadr2] = -deriv_2 else: # Single joint constraint - pos = d.qpos[worldid, qposadr1] - m.qpos0[qposadr1] - data[0] + pos = d.qpos[worldid, qposadr1] - m.qpos0[worldid, qposadr1] - data[0] Jqvel = d.qvel[worldid, dofadr1] - invweight = m.dof_invweight0[dofadr1] + invweight = m.dof_invweight0[worldid, dofadr1] # Update constraint parameters _update_efc_row( @@ -209,8 +211,8 @@ def _efc_equality_joint( pos, pos, invweight, - m.eq_solref[i_eq], - m.eq_solimp[i_eq], + m.eq_solref[worldid, i_eq], + m.eq_solimp[worldid, i_eq], wp.float32(0.0), Jqvel, 0.0, @@ -226,7 +228,7 @@ def _efc_friction( # TODO(team): tendon worldid, dofid = wp.tid() - if m.dof_frictionloss[dofid] <= 0.0: + if m.dof_frictionloss[worldid, dofid] <= 0.0: return efcid = wp.atomic_add(d.nefc, 0, 1) @@ -242,13 +244,13 @@ def _efc_friction( efcid, 0.0, 0.0, - m.dof_invweight0[dofid], - m.dof_solref[dofid], - m.dof_solimp[dofid], + m.dof_invweight0[worldid, dofid], + m.dof_solref[worldid, dofid], + m.dof_solimp[worldid, dofid], 0.0, Jqvel, - m.dof_frictionloss[dofid], - dofid, + m.dof_frictionloss[worldid, dofid], + m.jnt_bodyid[m.dof_jntid[dofid]], ) @@ -272,7 +274,7 @@ def _efc_equality_weld( obj1id = m.eq_obj1id[i_eq] obj2id = m.eq_obj2id[i_eq] - data = m.eq_data[i_eq] + data = m.eq_data[worldid, i_eq] 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]) @@ -284,9 +286,10 @@ def _efc_equality_weld( body2id = m.site_bodyid[obj2id] pos1 = d.site_xpos[worldid, obj1id] pos2 = d.site_xpos[worldid, obj2id] - - quat = math.mul_quat(d.xquat[worldid, body1id], m.site_quat[obj1id]) - quat1 = math.quat_inv(math.mul_quat(d.xquat[worldid, body2id], m.site_quat[obj2id])) + quat = math.mul_quat(d.xquat[worldid, body1id], m.site_quat[worldid, obj1id]) + quat1 = math.quat_inv( + math.mul_quat(d.xquat[worldid, body2id], m.site_quat[worldid, obj2id]) + ) else: body1id = obj1id @@ -325,12 +328,14 @@ def _efc_equality_weld( crotq = math.mul_quat(quat1, quat) # copy axis components crot = wp.vec3(crotq[1], crotq[2], crotq[3]) * torquescale - invweight_t = m.body_invweight0[body1id, 0] + m.body_invweight0[body2id, 0] + invweight_t = ( + m.body_invweight0[worldid, body1id, 0] + m.body_invweight0[worldid, body2id, 0] + ) pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) - solref = m.eq_solref[i_eq] - solimp = m.eq_solimp[i_eq] + solref = m.eq_solref[worldid, i_eq] + solimp = m.eq_solimp[worldid, i_eq] for i in range(3): _update_efc_row( @@ -348,7 +353,9 @@ def _efc_equality_weld( i_eq, ) - invweight_r = m.body_invweight0[body1id, 1] + m.body_invweight0[body2id, 1] + invweight_r = ( + m.body_invweight0[worldid, body1id, 1] + m.body_invweight0[worldid, body2id, 1] + ) for i in range(3): _update_efc_row( @@ -376,9 +383,9 @@ def _efc_limit_slide_hinge( jntid = m.jnt_limited_slide_hinge_adr[jntlimitedid] qpos = d.qpos[worldid, m.jnt_qposadr[jntid]] - jnt_range = m.jnt_range[jntid] + jnt_range = m.jnt_range[worldid, jntid] dist_min, dist_max = qpos - jnt_range[0], jnt_range[1] - qpos - pos = wp.min(dist_min, dist_max) - m.jnt_margin[jntid] + pos = wp.min(dist_min, dist_max) - m.jnt_margin[worldid, jntid] active = pos < 0 if active: @@ -398,10 +405,10 @@ def _efc_limit_slide_hinge( efcid, pos, pos, - m.dof_invweight0[dofadr], - m.jnt_solref[jntid], - m.jnt_solimp[jntid], - m.jnt_margin[jntid], + m.dof_invweight0[worldid, dofadr], + m.jnt_solref[worldid, jntid], + m.jnt_solimp[worldid, jntid], + m.jnt_margin[worldid, jntid], Jqvel, 0.0, dofadr, @@ -423,8 +430,8 @@ def _efc_limit_ball( ) axis_angle = math.quat_to_vel(jnt_quat) axis, angle = math.normalize_with_norm(axis_angle) - jnt_margin = m.jnt_margin[jntid] - jnt_range = m.jnt_range[jntid] + jnt_margin = m.jnt_margin[worldid, jntid] + jnt_range = m.jnt_range[worldid, jntid] pos = wp.max(jnt_range[0], jnt_range[1]) - angle - jnt_margin active = pos < 0 @@ -450,10 +457,10 @@ def _efc_limit_ball( efcid, pos, pos, - m.dof_invweight0[dofadr], - m.jnt_solref[jntid], - m.jnt_solimp[jntid], - jnt_margin, + m.dof_invweight0[worldid, dofadr], + m.jnt_solref[worldid, jntid], + m.jnt_solimp[worldid, jntid], + m.jnt_margin[worldid, jntid], Jqvel, 0.0, jntid, @@ -468,10 +475,10 @@ def _efc_limit_tendon( worldid, tenlimitedid = wp.tid() tenid = m.tendon_limited_adr[tenlimitedid] - ten_range = m.tendon_range[tenid] + ten_range = m.tendon_range[worldid, tenid] length = d.ten_length[worldid, tenid] dist_min, dist_max = length - ten_range[0], ten_range[1] - length - ten_margin = m.tendon_margin[tenid] + ten_margin = m.tendon_margin[worldid, tenid] pos = wp.min(dist_min, dist_max) - ten_margin active = pos < 0 @@ -503,9 +510,9 @@ def _efc_limit_tendon( efcid, pos, pos, - m.tendon_invweight0[tenid], - m.tendon_solref_lim[tenid], - m.tendon_solimp_lim[tenid], + m.tendon_invweight0[worldid, tenid], + m.tendon_solref_lim[worldid, tenid], + m.tendon_solimp_lim[worldid, tenid], ten_margin, Jqvel, 0.0, @@ -547,7 +554,9 @@ def _efc_contact_pyramidal( frame = d.contact.frame[conid] # pyramidal has common invweight across all edges - invweight = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0] + invweight = ( + m.body_invweight0[worldid, body1, 0] + m.body_invweight0[worldid, body2, 0] + ) if condim > 1: dimid2 = dimid / 2 + 1 @@ -648,7 +657,9 @@ def _efc_contact_elliptic( d.efc.J[efcid, i] = J Jqvel += J * d.qvel[worldid, i] - invweight = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0] + invweight = ( + m.body_invweight0[worldid, body1, 0] + m.body_invweight0[worldid, body2, 0] + ) ref = d.contact.solref[conid] pos_aref = pos diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index aaea9e87..76b4c703 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -128,7 +128,7 @@ def _next_act( # advance the actuation if m.actuator_dyntype[actid] == wp.static(DynType.FILTEREXACT.value): - dyn_prm = m.actuator_dynprm[actid] + dyn_prm = m.actuator_dynprm[worldid, actid] tau = wp.max(MJ_MINVAL, dyn_prm[0]) act += act_dot * tau * (1.0 - wp.exp(-m.opt.timestep / tau)) else: @@ -136,7 +136,7 @@ def _next_act( # clamp to actrange if m.actuator_actlimited[actid]: - actrange = m.actuator_actrange[actid] + actrange = m.actuator_actrange[worldid, actid] act = wp.clamp(act, actrange[0], actrange[1]) d.act[worldid, actid] = act @@ -193,7 +193,9 @@ def add_damping_sum_qfrc_kernel_sparse(m: Model, d: Data): worldid, tid = wp.tid() dof_Madr = m.dof_Madr[tid] - d.qM_integration[worldid, 0, dof_Madr] += m.opt.timestep * m.dof_damping[tid] + d.qM_integration[worldid, 0, dof_Madr] += ( + m.opt.timestep * m.dof_damping[worldid, tid] + ) d.qfrc_integration[worldid, tid] = ( d.qfrc_smooth[worldid, tid] + d.qfrc_constraint[worldid, tid] @@ -215,14 +217,16 @@ def eulerdamp_fused_dense(m: Model, d: Data): def tile_eulerdamp(adr: int, size: int, tilesize: int): @kernel def eulerdamp( - m: Model, d: Data, damping: wp.array(dtype=wp.float32), leveladr: int + m: Model, d: Data, damping: wp.array2d(dtype=wp.float32), leveladr: int ): worldid, nodeid = wp.tid() dofid = m.qLD_tile[leveladr + nodeid] M_tile = wp.tile_load( d.qM[worldid], shape=(tilesize, tilesize), offset=(dofid, dofid) ) - damping_tile = wp.tile_load(damping, shape=(tilesize,), offset=(dofid,)) + damping_tile = wp.tile_load( + damping[worldid], shape=(tilesize,), offset=(dofid,) + ) damping_scaled = damping_tile * m.opt.timestep qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled) @@ -383,10 +387,10 @@ def actuator_bias_gain_vel(m: Model, d: Data): actuator_dyntype = m.actuator_dyntype[actid] if actuator_biastype == wp.static(BiasType.AFFINE.value): - bias_vel = m.actuator_biasprm[actid][2] + bias_vel = m.actuator_biasprm[worldid, actid][2] if actuator_gaintype == wp.static(GainType.AFFINE.value): - gain_vel = m.actuator_gainprm[actid][2] + gain_vel = m.actuator_gainprm[worldid, actid][2] ctrl = d.ctrl[worldid, actid] @@ -396,7 +400,7 @@ def actuator_bias_gain_vel(m: Model, d: Data): d.act_vel_integration[worldid, actid] = bias_vel + gain_vel * ctrl def qderiv_actuator_damping_fused( - m: Model, d: Data, damping: wp.array(dtype=wp.float32) + m: Model, d: Data, damping: wp.array2d(dtype=wp.float32) ): if actuation_enabled: block_dim = 64 @@ -412,7 +416,7 @@ def qderiv_actuator_damping_tiled( ): @kernel def qderiv_actuator_fused_kernel( - m: Model, d: Data, damping: wp.array(dtype=wp.float32), leveladr: int + m: Model, d: Data, damping: wp.array2d(dtype=wp.float32), leveladr: int ): worldid, nodeid = wp.tid() offset_nv = m.actuator_moment_offset_nv[leveladr + nodeid] @@ -439,7 +443,9 @@ def qderiv_actuator_fused_kernel( ) if wp.static(passive_enabled): - dof_damping = wp.tile_load(damping, shape=tilesize_nv, offset=offset_nv) + dof_damping = wp.tile_load( + damping[worldid], shape=tilesize_nv, offset=offset_nv + ) negative = wp.neg(dof_damping) qderiv_tile = wp.tile_diag_add(qderiv_tile, negative) @@ -632,7 +638,7 @@ def _force(m: Model, d: Data): dsbl_clampctrl = m.opt.disableflags & wp.static(DisableBit.CLAMPCTRL.value) if m.actuator_ctrllimited[uid] and not dsbl_clampctrl: - r = m.actuator_ctrlrange[uid] + r = m.actuator_ctrlrange[worldid, uid] ctrl = wp.clamp(ctrl, r[0], r[1]) if m.na: @@ -643,7 +649,7 @@ def _force(m: Model, d: Data): elif dyntype == int(DynType.FILTER.value) or dyntype == int( DynType.FILTEREXACT.value ): - dynprm = m.actuator_dynprm[uid] + dynprm = m.actuator_dynprm[worldid, uid] actadr = m.actuator_actadr[uid] act = d.act[worldid, actadr] d.act_dot[worldid, actadr] = (ctrl - act) / wp.max(dynprm[0], MJ_MINVAL) @@ -662,7 +668,7 @@ def _force(m: Model, d: Data): # gain gaintype = m.actuator_gaintype[uid] - gainprm = m.actuator_gainprm[uid] + gainprm = m.actuator_gainprm[worldid, uid] gain = 0.0 if gaintype == int(GainType.FIXED.value): @@ -674,7 +680,7 @@ def _force(m: Model, d: Data): # bias biastype = m.actuator_biastype[uid] - biasprm = m.actuator_biasprm[uid] + biasprm = m.actuator_biasprm[worldid, uid] bias = 0.0 # BiasType.NONE if biastype == int(BiasType.AFFINE.value): @@ -687,7 +693,7 @@ def _force(m: Model, d: Data): # TODO(team): tendon total force clamping if m.actuator_forcelimited[uid]: - r = m.actuator_forcerange[uid] + r = m.actuator_forcerange[worldid, uid] f = wp.clamp(f, r[0], r[1]) d.actuator_force[worldid, uid] = f @@ -696,10 +702,11 @@ def _qfrc_limited(m: Model, d: Data): worldid, dofid = wp.tid() jntid = m.dof_jntid[dofid] if m.jnt_actfrclimited[jntid]: + range = m.jnt_actfrcrange[worldid, jntid] d.qfrc_actuator[worldid, dofid] = wp.clamp( d.qfrc_actuator[worldid, dofid], - m.jnt_actfrcrange[jntid][0], - m.jnt_actfrcrange[jntid][1], + range[0], + range[1], ) if m.opt.is_sparse: @@ -714,8 +721,8 @@ def _qfrc(m: Model, moment: array3df, force: array2df, qfrc: array2df): s += moment[worldid, uid, vid] * force[worldid, uid] jntid = m.dof_jntid[vid] if m.jnt_actfrclimited[jntid]: - r = m.jnt_actfrcrange[jntid] - s = wp.clamp(s, r[0], r[1]) + range = m.jnt_actfrcrange[worldid, jntid] + s = wp.clamp(s, range[0], range[1]) qfrc[worldid, vid] = s wp.launch(_force, dim=[d.nworld, m.nu], inputs=[m, d]) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 0a1638c1..b29bb1d2 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -70,6 +70,14 @@ def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]: return np.array(geompairs), np.array(pairids) +def create_nmodel_batched_array(mjm_array, dtype): + array = wp.array(mjm_array, dtype=dtype) + array.ndim += 1 + array.shape = (1,) + array.shape + array.strides = (0,) + array.strides + return array + + def put_model(mjm: mujoco.MjModel) -> types.Model: # check supported features for field, field_types, field_str in ( @@ -158,8 +166,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.opt.depth_extension = wp.float32(0.1) # warp only m.stat.meaninertia = mjm.stat.meaninertia - m.qpos0 = wp.array(mjm.qpos0, dtype=wp.float32, ndim=1) - m.qpos_spring = wp.array(mjm.qpos_spring, dtype=wp.float32, ndim=1) + m.qpos0 = create_nmodel_batched_array(mjm.qpos0, dtype=wp.float32) + m.qpos_spring = create_nmodel_batched_array(mjm.qpos_spring, dtype=wp.float32) # dof lower triangle row and column indices dof_tri_row, dof_tri_col = np.tril_indices(mjm.nv) @@ -360,22 +368,24 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.body_parentid = wp.array(mjm.body_parentid, dtype=wp.int32, ndim=1) m.body_mocapid = wp.array(mjm.body_mocapid, dtype=wp.int32, ndim=1) m.body_weldid = wp.array(mjm.body_weldid, dtype=wp.int32, ndim=1) - m.body_pos = wp.array(mjm.body_pos, dtype=wp.vec3, ndim=1) - m.body_quat = wp.array(mjm.body_quat, dtype=wp.quat, ndim=1) - m.body_ipos = wp.array(mjm.body_ipos, dtype=wp.vec3, ndim=1) - m.body_iquat = wp.array(mjm.body_iquat, dtype=wp.quat, ndim=1) + m.body_pos = create_nmodel_batched_array(mjm.body_pos, dtype=wp.vec3) + m.body_quat = create_nmodel_batched_array(mjm.body_quat, dtype=wp.quat) + m.body_ipos = create_nmodel_batched_array(mjm.body_ipos, dtype=wp.vec3) + m.body_iquat = create_nmodel_batched_array(mjm.body_iquat, dtype=wp.quat) m.body_rootid = wp.array(mjm.body_rootid, dtype=wp.int32, ndim=1) - m.body_inertia = wp.array(mjm.body_inertia, dtype=wp.vec3, ndim=1) - m.body_mass = wp.array(mjm.body_mass, dtype=wp.float32, ndim=1) - m.body_subtreemass = wp.array(mjm.body_subtreemass, dtype=wp.float32, ndim=1) + m.body_inertia = create_nmodel_batched_array(mjm.body_inertia, dtype=wp.vec3) + m.body_mass = create_nmodel_batched_array(mjm.body_mass, dtype=wp.float32) + m.body_subtreemass = create_nmodel_batched_array( + mjm.body_subtreemass, dtype=wp.float32 + ) subtree_mass = np.copy(mjm.body_mass) # TODO(team): should this be [mjm.nbody - 1, 0) ? for i in range(mjm.nbody - 1, -1, -1): subtree_mass[mjm.body_parentid[i]] += subtree_mass[i] - m.subtree_mass = wp.array(subtree_mass, dtype=wp.float32, ndim=1) - m.body_invweight0 = wp.array(mjm.body_invweight0, dtype=wp.float32, ndim=2) + m.subtree_mass = create_nmodel_batched_array(subtree_mass, dtype=wp.float32) + m.body_invweight0 = create_nmodel_batched_array(mjm.body_invweight0, dtype=wp.float32) m.body_geomnum = wp.array(mjm.body_geomnum, dtype=wp.int32, ndim=1) m.body_geomadr = wp.array(mjm.body_geomadr, dtype=wp.int32, ndim=1) m.body_contype = wp.array(mjm.body_contype, dtype=wp.int32, ndim=1) @@ -387,35 +397,35 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: ) m.jnt_limited_ball_adr = wp.array(jnt_limited_ball_adr, dtype=wp.int32, ndim=1) m.jnt_type = wp.array(mjm.jnt_type, dtype=wp.int32, ndim=1) - m.jnt_solref = wp.array(mjm.jnt_solref, dtype=wp.vec2f, ndim=1) - m.jnt_solimp = wp.array(mjm.jnt_solimp, dtype=types.vec5, ndim=1) + m.jnt_solref = create_nmodel_batched_array(mjm.jnt_solref, dtype=wp.vec2f) + m.jnt_solimp = create_nmodel_batched_array(mjm.jnt_solimp, dtype=types.vec5) m.jnt_qposadr = wp.array(mjm.jnt_qposadr, dtype=wp.int32, ndim=1) m.jnt_dofadr = wp.array(mjm.jnt_dofadr, dtype=wp.int32, ndim=1) m.jnt_axis = wp.array(mjm.jnt_axis, dtype=wp.vec3, ndim=1) - m.jnt_pos = wp.array(mjm.jnt_pos, dtype=wp.vec3, ndim=1) - m.jnt_range = wp.array(mjm.jnt_range, dtype=wp.float32, ndim=2) - m.jnt_margin = wp.array(mjm.jnt_margin, dtype=wp.float32, ndim=1) - m.jnt_stiffness = wp.array(mjm.jnt_stiffness, dtype=wp.float32, ndim=1) + m.jnt_pos = create_nmodel_batched_array(mjm.jnt_pos, dtype=wp.vec3) + m.jnt_range = create_nmodel_batched_array(mjm.jnt_range, dtype=wp.float32) + m.jnt_stiffness = create_nmodel_batched_array(mjm.jnt_stiffness, dtype=wp.float32) + m.jnt_margin = create_nmodel_batched_array(mjm.jnt_margin, dtype=wp.float32) m.jnt_actfrclimited = wp.array(mjm.jnt_actfrclimited, dtype=wp.bool, ndim=1) - m.jnt_actfrcrange = wp.array(mjm.jnt_actfrcrange, dtype=wp.vec2, ndim=1) + m.jnt_actfrcrange = create_nmodel_batched_array(mjm.jnt_actfrcrange, dtype=wp.vec2) m.geom_type = wp.array(mjm.geom_type, dtype=wp.int32, ndim=1) m.geom_bodyid = wp.array(mjm.geom_bodyid, dtype=wp.int32, ndim=1) - m.geom_conaffinity = wp.array(mjm.geom_conaffinity, dtype=wp.int32, ndim=1) - m.geom_contype = wp.array(mjm.geom_contype, dtype=wp.int32, ndim=1) + m.geom_conaffinity = create_nmodel_batched_array(mjm.geom_conaffinity, dtype=wp.int32) + m.geom_contype = create_nmodel_batched_array(mjm.geom_contype, dtype=wp.int32) m.geom_condim = wp.array(mjm.geom_condim, dtype=wp.int32, ndim=1) - m.geom_pos = wp.array(mjm.geom_pos, dtype=wp.vec3, ndim=1) - m.geom_quat = wp.array(mjm.geom_quat, dtype=wp.quat, ndim=1) - m.geom_size = wp.array(mjm.geom_size, dtype=wp.vec3, ndim=1) - m.geom_priority = wp.array(mjm.geom_priority, dtype=wp.int32, ndim=1) - m.geom_solmix = wp.array(mjm.geom_solmix, dtype=wp.float32, ndim=1) - m.geom_solref = wp.array(mjm.geom_solref, dtype=wp.vec2, ndim=1) - m.geom_solimp = wp.array(mjm.geom_solimp, dtype=types.vec5, ndim=1) - m.geom_friction = wp.array(mjm.geom_friction, dtype=wp.vec3, ndim=1) - m.geom_margin = wp.array(mjm.geom_margin, dtype=wp.float32, ndim=1) - m.geom_gap = wp.array(mjm.geom_gap, dtype=wp.float32, ndim=1) + m.geom_pos = create_nmodel_batched_array(mjm.geom_pos, dtype=wp.vec3) + m.geom_quat = create_nmodel_batched_array(mjm.geom_quat, dtype=wp.quat) + m.geom_size = create_nmodel_batched_array(mjm.geom_size, dtype=wp.vec3) + m.geom_priority = create_nmodel_batched_array(mjm.geom_priority, dtype=wp.int32) + m.geom_solmix = create_nmodel_batched_array(mjm.geom_solmix, dtype=wp.float32) + m.geom_solref = create_nmodel_batched_array(mjm.geom_solref, dtype=wp.vec2) + m.geom_solimp = create_nmodel_batched_array(mjm.geom_solimp, dtype=types.vec5) + m.geom_friction = create_nmodel_batched_array(mjm.geom_friction, dtype=wp.vec3) + m.geom_margin = create_nmodel_batched_array(mjm.geom_margin, dtype=wp.float32) + m.geom_gap = create_nmodel_batched_array(mjm.geom_gap, dtype=wp.float32) m.geom_aabb = wp.array(mjm.geom_aabb, dtype=wp.vec3, ndim=3) - m.geom_rbound = wp.array(mjm.geom_rbound, dtype=wp.float32, ndim=1) - m.geom_dataid = wp.array(mjm.geom_dataid, dtype=wp.int32, ndim=1) + m.geom_rbound = create_nmodel_batched_array(mjm.geom_rbound, dtype=wp.float32) + m.geom_dataid = create_nmodel_batched_array(mjm.geom_dataid, dtype=wp.int32) m.mesh_vertadr = wp.array(mjm.mesh_vertadr, dtype=wp.int32, ndim=1) m.mesh_vertnum = wp.array(mjm.mesh_vertnum, dtype=wp.int32, ndim=1) m.mesh_vert = wp.array(mjm.mesh_vert, dtype=wp.vec3, ndim=1) @@ -424,55 +434,71 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.eq_obj2id = wp.array(mjm.eq_obj2id, dtype=wp.int32, ndim=1) m.eq_objtype = wp.array(mjm.eq_objtype, dtype=wp.int32, ndim=1) m.eq_active0 = wp.array(mjm.eq_active0, dtype=wp.bool, ndim=1) - m.eq_solref = wp.array(mjm.eq_solref, dtype=wp.vec2, ndim=1) - m.eq_solimp = wp.array(mjm.eq_solimp, dtype=types.vec5, ndim=1) - m.eq_data = wp.array(mjm.eq_data, dtype=types.vec11, ndim=1) - m.site_pos = wp.array(mjm.site_pos, dtype=wp.vec3, ndim=1) - m.site_quat = wp.array(mjm.site_quat, dtype=wp.quat, ndim=1) + m.eq_solref = create_nmodel_batched_array(mjm.eq_solref, dtype=wp.vec2) + m.eq_solimp = create_nmodel_batched_array(mjm.eq_solimp, dtype=types.vec5) + m.eq_data = create_nmodel_batched_array(mjm.eq_data, dtype=types.vec11) + m.site_pos = create_nmodel_batched_array(mjm.site_pos, dtype=wp.vec3) + m.site_quat = create_nmodel_batched_array(mjm.site_quat, dtype=wp.quat) m.site_bodyid = wp.array(mjm.site_bodyid, dtype=wp.int32, ndim=1) m.cam_mode = wp.array(mjm.cam_mode, dtype=wp.int32, ndim=1) m.cam_bodyid = wp.array(mjm.cam_bodyid, dtype=wp.int32, ndim=1) m.cam_targetbodyid = wp.array(mjm.cam_targetbodyid, dtype=wp.int32, ndim=1) - m.cam_pos = wp.array(mjm.cam_pos, dtype=wp.vec3, ndim=1) - m.cam_quat = wp.array(mjm.cam_quat, dtype=wp.quat, ndim=1) - m.cam_poscom0 = wp.array(mjm.cam_poscom0, dtype=wp.vec3, ndim=1) - m.cam_pos0 = wp.array(mjm.cam_pos0, dtype=wp.vec3, ndim=1) + m.cam_pos = create_nmodel_batched_array(mjm.cam_pos, dtype=wp.vec3) + m.cam_quat = create_nmodel_batched_array(mjm.cam_quat, dtype=wp.quat) + m.cam_poscom0 = create_nmodel_batched_array(mjm.cam_poscom0, dtype=wp.vec3) + m.cam_pos0 = create_nmodel_batched_array(mjm.cam_pos0, dtype=wp.vec3) m.light_mode = wp.array(mjm.light_mode, dtype=wp.int32, ndim=1) m.light_bodyid = wp.array(mjm.light_bodyid, dtype=wp.int32, ndim=1) m.light_targetbodyid = wp.array(mjm.light_targetbodyid, dtype=wp.int32, ndim=1) - m.light_pos = wp.array(mjm.light_pos, dtype=wp.vec3, ndim=1) - m.light_dir = wp.array(mjm.light_dir, dtype=wp.vec3, ndim=1) - m.light_poscom0 = wp.array(mjm.light_poscom0, dtype=wp.vec3, ndim=1) - m.light_pos0 = wp.array(mjm.light_pos0, dtype=wp.vec3, ndim=1) + m.light_pos = create_nmodel_batched_array(mjm.light_pos, dtype=wp.vec3) + m.light_dir = create_nmodel_batched_array(mjm.light_dir, dtype=wp.vec3) + m.light_poscom0 = create_nmodel_batched_array(mjm.light_poscom0, dtype=wp.vec3) + m.light_pos0 = create_nmodel_batched_array(mjm.light_pos0, dtype=wp.vec3) m.dof_bodyid = wp.array(mjm.dof_bodyid, dtype=wp.int32, ndim=1) m.dof_jntid = wp.array(mjm.dof_jntid, dtype=wp.int32, ndim=1) m.dof_parentid = wp.array(mjm.dof_parentid, dtype=wp.int32, ndim=1) m.dof_Madr = wp.array(mjm.dof_Madr, dtype=wp.int32, ndim=1) - m.dof_armature = wp.array(mjm.dof_armature, dtype=wp.float32, ndim=1) - m.dof_damping = wp.array(mjm.dof_damping, dtype=wp.float32, ndim=1) - m.dof_frictionloss = wp.array(mjm.dof_frictionloss, dtype=wp.float32, ndim=1) - m.dof_solimp = wp.array(mjm.dof_solimp, dtype=types.vec5, ndim=1) - m.dof_solref = wp.array(mjm.dof_solref, dtype=wp.vec2, ndim=1) + m.dof_armature = create_nmodel_batched_array(mjm.dof_armature, dtype=wp.float32) + m.dof_damping = create_nmodel_batched_array(mjm.dof_damping, dtype=wp.float32) + m.dof_frictionloss = create_nmodel_batched_array( + mjm.dof_frictionloss, dtype=wp.float32 + ) + m.dof_solimp = create_nmodel_batched_array(mjm.dof_solimp, dtype=types.vec5) + m.dof_solref = create_nmodel_batched_array(mjm.dof_solref, dtype=wp.vec2) m.dof_tri_row = wp.from_numpy(dof_tri_row, dtype=wp.int32) m.dof_tri_col = wp.from_numpy(dof_tri_col, dtype=wp.int32) - m.dof_invweight0 = wp.array(mjm.dof_invweight0, dtype=wp.float32, ndim=1) + m.dof_invweight0 = create_nmodel_batched_array(mjm.dof_invweight0, dtype=wp.float32) m.actuator_trntype = wp.array(mjm.actuator_trntype, dtype=wp.int32, ndim=1) m.actuator_trnid = wp.array(mjm.actuator_trnid, dtype=wp.int32, ndim=2) m.actuator_ctrllimited = wp.array(mjm.actuator_ctrllimited, dtype=wp.bool, ndim=1) - m.actuator_ctrlrange = wp.array(mjm.actuator_ctrlrange, dtype=wp.vec2, ndim=1) + m.actuator_ctrlrange = create_nmodel_batched_array( + mjm.actuator_ctrlrange, dtype=wp.vec2 + ) m.actuator_forcelimited = wp.array(mjm.actuator_forcelimited, dtype=wp.bool, ndim=1) - m.actuator_forcerange = wp.array(mjm.actuator_forcerange, dtype=wp.vec2, ndim=1) + m.actuator_forcerange = create_nmodel_batched_array( + mjm.actuator_forcerange, dtype=wp.vec2 + ) m.actuator_gaintype = wp.array(mjm.actuator_gaintype, dtype=wp.int32, ndim=1) - m.actuator_gainprm = wp.array(mjm.actuator_gainprm, dtype=types.vec10f, ndim=1) + m.actuator_gainprm = create_nmodel_batched_array( + mjm.actuator_gainprm, dtype=types.vec10f + ) m.actuator_biastype = wp.array(mjm.actuator_biastype, dtype=wp.int32, ndim=1) - m.actuator_biasprm = wp.array(mjm.actuator_biasprm, dtype=types.vec10f, ndim=1) - m.actuator_gear = wp.array(mjm.actuator_gear, dtype=wp.spatial_vector, ndim=1) + m.actuator_biasprm = create_nmodel_batched_array( + mjm.actuator_biasprm, dtype=types.vec10f + ) + m.actuator_gear = create_nmodel_batched_array( + mjm.actuator_gear, dtype=wp.spatial_vector + ) m.actuator_actlimited = wp.array(mjm.actuator_actlimited, dtype=wp.bool, ndim=1) - m.actuator_actrange = wp.array(mjm.actuator_actrange, dtype=wp.vec2, ndim=1) + m.actuator_actrange = create_nmodel_batched_array( + mjm.actuator_actrange, dtype=wp.vec2 + ) m.actuator_actadr = wp.array(mjm.actuator_actadr, dtype=wp.int32, ndim=1) m.actuator_actnum = wp.array(mjm.actuator_actnum, dtype=wp.int32, ndim=1) m.actuator_dyntype = wp.array(mjm.actuator_dyntype, dtype=wp.int32, ndim=1) - m.actuator_dynprm = wp.array(mjm.actuator_dynprm, dtype=types.vec10f, ndim=1) + m.actuator_dynprm = create_nmodel_batched_array( + mjm.actuator_dynprm, dtype=types.vec10f + ) m.exclude_signature = wp.array(mjm.exclude_signature, dtype=wp.int32, ndim=1) # pre-compute indices of equality constraints @@ -500,12 +526,14 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.pair_dim = wp.array(mjm.pair_dim, dtype=wp.int32, ndim=1) m.pair_geom1 = wp.array(mjm.pair_geom1, dtype=wp.int32, ndim=1) m.pair_geom2 = wp.array(mjm.pair_geom2, dtype=wp.int32, ndim=1) - m.pair_solref = wp.array(mjm.pair_solref, dtype=wp.vec2, ndim=1) - m.pair_solreffriction = wp.array(mjm.pair_solreffriction, dtype=wp.vec2, ndim=1) - m.pair_solimp = wp.array(mjm.pair_solimp, dtype=types.vec5, ndim=1) - m.pair_margin = wp.array(mjm.pair_margin, dtype=wp.float32, ndim=1) - m.pair_gap = wp.array(mjm.pair_gap, dtype=wp.float32, ndim=1) - m.pair_friction = wp.array(mjm.pair_friction, dtype=types.vec5, ndim=1) + m.pair_solref = create_nmodel_batched_array(mjm.pair_solref, dtype=wp.vec2) + m.pair_solreffriction = create_nmodel_batched_array( + mjm.pair_solreffriction, dtype=wp.vec2 + ) + m.pair_solimp = create_nmodel_batched_array(mjm.pair_solimp, dtype=types.vec5) + m.pair_margin = create_nmodel_batched_array(mjm.pair_margin, dtype=wp.float32) + m.pair_gap = create_nmodel_batched_array(mjm.pair_gap, dtype=wp.float32) + m.pair_friction = create_nmodel_batched_array(mjm.pair_friction, dtype=types.vec5) m.condim_max = np.max(mjm.geom_condim) # TODO(team): get max after filtering # tendon @@ -515,15 +543,21 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: m.tendon_limited_adr = wp.array( np.nonzero(mjm.tendon_limited)[0], dtype=wp.int32, ndim=1 ) - m.tendon_solref_lim = wp.array(mjm.tendon_solref_lim, dtype=wp.vec2f, ndim=1) - m.tendon_solimp_lim = wp.array(mjm.tendon_solimp_lim, dtype=types.vec5, ndim=1) - m.tendon_range = wp.array(mjm.tendon_range, dtype=wp.vec2f, ndim=1) - m.tendon_margin = wp.array(mjm.tendon_margin, dtype=wp.float32, ndim=1) - m.tendon_length0 = wp.array(mjm.tendon_length0, dtype=wp.float32, ndim=1) - m.tendon_invweight0 = wp.array(mjm.tendon_invweight0, dtype=wp.float32, ndim=1) - m.wrap_objid = wp.array(mjm.wrap_objid, dtype=wp.int32, ndim=1) - m.wrap_prm = wp.array(mjm.wrap_prm, dtype=wp.float32, ndim=1) - m.wrap_type = wp.array(mjm.wrap_type, dtype=wp.int32, ndim=1) + m.tendon_solref_lim = create_nmodel_batched_array( + mjm.tendon_solref_lim, dtype=wp.vec2f + ) + m.tendon_solimp_lim = create_nmodel_batched_array( + mjm.tendon_solimp_lim, dtype=types.vec5 + ) + m.tendon_range = create_nmodel_batched_array(mjm.tendon_range, dtype=wp.vec2f) + m.tendon_margin = create_nmodel_batched_array(mjm.tendon_margin, dtype=wp.float32) + m.tendon_length0 = create_nmodel_batched_array(mjm.tendon_length0, dtype=wp.float32) + m.tendon_invweight0 = create_nmodel_batched_array( + mjm.tendon_invweight0, dtype=wp.float32 + ) + m.wrap_objid = wp.array(mjm.wrap_objid, dtype=wp.int32) + m.wrap_prm = create_nmodel_batched_array(mjm.wrap_prm, dtype=wp.float32) + m.wrap_type = wp.array(mjm.wrap_type, dtype=wp.int32) # fixed tendon tendon_jnt_adr = [] diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index c54f8bc7..72d02f34 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -24,6 +24,16 @@ from . import test_util +# tolerance for difference between MuJoCo and MJWarp smooth calculations - mostly +# due to float precision +_TOLERANCE = 5e-5 + + +def _assert_eq(a, b, name): + tol = _TOLERANCE * 10 # avoid test noise + err_msg = f"mismatch: {name}" + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) + class IOTest(absltest.TestCase): def test_make_put_data(self): @@ -290,6 +300,40 @@ def test_option_physical_constants(self): with self.assertRaises(NotImplementedError): mjwarp.put_model(mjm) + def test_model_batching(self): + mjm, mjd, _, _ = test_util.fixture("humanoid/humanoid.xml", kick=True) + + m = mjwarp.put_model(mjm) + d = mjwarp.put_data(mjm, mjd, nworld=3) + + # manually create a batch of damping values + damping_orig = mjm.dof_damping + dof_damping = np.zeros((2, len(damping_orig)), dtype=np.float32) + dof_damping[0, :] = damping_orig + dof_damping[1, :] = damping_orig * 0.5 + + # set the batched damping values + m.dof_damping = wp.from_numpy(dof_damping, dtype=wp.float32) + + mjwarp.passive(m, d) + + # mujoco reference, just have 3 separate model/data strctures + mujoco.mj_passive(mjm, mjd) + + m2, d2, _, _ = test_util.fixture("humanoid/humanoid.xml") + d2.qvel = mjd.qvel # need to copy qvel because of randomization + m2.dof_damping *= 0.5 + + m3, d3, _, _ = test_util.fixture("humanoid/humanoid.xml") + d3.qvel = mjd.qvel # need to copy qvel because of randomization + + mujoco.mj_passive(m2, d2) + mujoco.mj_passive(m3, d3) + + _assert_eq(d.qfrc_damper.numpy()[0, :], mjd.qfrc_damper, "qfrc_damper") + _assert_eq(d.qfrc_damper.numpy()[1, :], d2.qfrc_damper, "qfrc_damper") + _assert_eq(d.qfrc_damper.numpy()[0, :], d3.qfrc_damper, "qfrc_damper") + if __name__ == "__main__": wp.init() diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index c8cbc30b..033aad9a 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -35,7 +35,7 @@ def passive(m: Model, d: Data): @kernel def _spring(m: Model, d: Data): worldid, jntid = wp.tid() - stiffness = m.jnt_stiffness[jntid] + stiffness = m.jnt_stiffness[worldid, jntid] dofid = m.jnt_dofadr[jntid] if stiffness == 0.0: @@ -46,9 +46,9 @@ def _spring(m: Model, d: Data): if jnt_type == wp.static(JointType.FREE.value): dif = wp.vec3( - d.qpos[worldid, qposid + 0] - m.qpos_spring[qposid + 0], - d.qpos[worldid, qposid + 1] - m.qpos_spring[qposid + 1], - d.qpos[worldid, qposid + 2] - m.qpos_spring[qposid + 2], + d.qpos[worldid, qposid + 0] - m.qpos_spring[worldid, qposid + 0], + d.qpos[worldid, qposid + 1] - m.qpos_spring[worldid, qposid + 1], + d.qpos[worldid, qposid + 2] - m.qpos_spring[worldid, qposid + 2], ) d.qfrc_spring[worldid, dofid + 0] = -stiffness * dif[0] d.qfrc_spring[worldid, dofid + 1] = -stiffness * dif[1] @@ -60,10 +60,10 @@ def _spring(m: Model, d: Data): d.qpos[worldid, qposid + 6], ) ref = wp.quat( - m.qpos_spring[qposid + 3], - m.qpos_spring[qposid + 4], - m.qpos_spring[qposid + 5], - m.qpos_spring[qposid + 6], + m.qpos_spring[worldid, qposid + 3], + m.qpos_spring[worldid, qposid + 4], + m.qpos_spring[worldid, qposid + 5], + m.qpos_spring[worldid, qposid + 6], ) dif = math.quat_sub(rot, ref) d.qfrc_spring[worldid, dofid + 3] = -stiffness * dif[0] @@ -77,23 +77,23 @@ def _spring(m: Model, d: Data): d.qpos[worldid, qposid + 3], ) ref = wp.quat( - m.qpos_spring[qposid + 0], - m.qpos_spring[qposid + 1], - m.qpos_spring[qposid + 2], - m.qpos_spring[qposid + 3], + m.qpos_spring[worldid, qposid + 0], + m.qpos_spring[worldid, qposid + 1], + m.qpos_spring[worldid, qposid + 2], + m.qpos_spring[worldid, qposid + 3], ) dif = math.quat_sub(rot, ref) d.qfrc_spring[worldid, dofid + 0] = -stiffness * dif[0] d.qfrc_spring[worldid, dofid + 1] = -stiffness * dif[1] d.qfrc_spring[worldid, dofid + 2] = -stiffness * dif[2] else: # mjJNT_SLIDE, mjJNT_HINGE - fdif = d.qpos[worldid, qposid] - m.qpos_spring[qposid] + fdif = d.qpos[worldid, qposid] - m.qpos_spring[worldid, qposid] d.qfrc_spring[worldid, dofid] = -stiffness * fdif @kernel def _damper_passive(m: Model, d: Data): worldid, dofid = wp.tid() - damping = m.dof_damping[dofid] + damping = m.dof_damping[worldid, dofid] qfrc_damper = -damping * d.qvel[worldid, dofid] d.qfrc_damper[worldid, dofid] = qfrc_damper diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 1e55a0e1..ae02738a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -167,25 +167,37 @@ def _frame_quat( m: Model, d: Data, worldid: int, objid: int, objtype: int, refid: int ) -> wp.quat: if objtype == int(ObjType.BODY.value): - quat = math.mul_quat(d.xquat[worldid, objid], m.body_iquat[objid]) + quat = math.mul_quat(d.xquat[worldid, objid], m.body_iquat[worldid, objid]) if refid == -1: return quat - refquat = math.mul_quat(d.xquat[worldid, refid], m.body_iquat[refid]) + refquat = math.mul_quat(d.xquat[worldid, refid], m.body_iquat[worldid, refid]) elif objtype == int(ObjType.XBODY.value): quat = d.xquat[worldid, objid] if refid == -1: return quat - refquat = d.xquat[worldid, refid] + refquat = math.mul_quat(d.xquat[worldid, refid], m.body_iquat[worldid, refid]) elif objtype == int(ObjType.GEOM.value): - quat = math.mul_quat(d.xquat[worldid, m.geom_bodyid[objid]], m.geom_quat[objid]) + quat = math.mul_quat( + d.xquat[worldid, m.geom_bodyid[objid]], + m.geom_quat[worldid, objid], + ) if refid == -1: return quat - refquat = math.mul_quat(d.xquat[worldid, m.geom_bodyid[refid]], m.geom_quat[refid]) + refquat = math.mul_quat( + d.xquat[worldid, m.geom_bodyid[refid]], + m.geom_quat[worldid, refid], + ) elif objtype == int(ObjType.SITE.value): - quat = math.mul_quat(d.xquat[worldid, m.site_bodyid[objid]], m.site_quat[objid]) + quat = math.mul_quat( + d.xquat[worldid, m.site_bodyid[objid]], + m.site_quat[worldid, objid], + ) if refid == -1: return quat - refquat = math.mul_quat(d.xquat[worldid, m.site_bodyid[refid]], m.site_quat[refid]) + refquat = math.mul_quat( + d.xquat[worldid, m.site_bodyid[refid]], + m.site_quat[worldid, refid], + ) # TODO(team): camera diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index bc6a0ae6..ff7d4c62 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -61,8 +61,8 @@ def _level(m: Model, d: Data, leveladr: int): if jntnum == 0: # no joints - apply fixed translation and rotation relative to parent pid = m.body_parentid[bodyid] - xpos = (d.xmat[worldid, pid] * m.body_pos[bodyid]) + d.xpos[worldid, pid] - xquat = math.mul_quat(d.xquat[worldid, pid], m.body_quat[bodyid]) + xpos = (d.xmat[worldid, pid] * m.body_pos[worldid, bodyid]) + d.xpos[worldid, pid] + xquat = math.mul_quat(d.xquat[worldid, pid], m.body_quat[worldid, bodyid]) elif jntnum == 1 and m.jnt_type[jntadr] == wp.static(JointType.FREE.value): # free joint qadr = m.jnt_qposadr[jntadr] @@ -74,14 +74,14 @@ def _level(m: Model, d: Data, leveladr: int): # regular or no joints # apply fixed translation and rotation relative to parent pid = m.body_parentid[bodyid] - xpos = (d.xmat[worldid, pid] * m.body_pos[bodyid]) + d.xpos[worldid, pid] - xquat = math.mul_quat(d.xquat[worldid, pid], m.body_quat[bodyid]) + xpos = (d.xmat[worldid, pid] * m.body_pos[worldid, bodyid]) + d.xpos[worldid, pid] + xquat = math.mul_quat(d.xquat[worldid, pid], m.body_quat[worldid, bodyid]) for _ in range(jntnum): qadr = m.jnt_qposadr[jntadr] jnt_type = m.jnt_type[jntadr] jnt_axis = m.jnt_axis[jntadr] - xanchor = math.rot_vec_quat(m.jnt_pos[jntadr], xquat) + xpos + xanchor = math.rot_vec_quat(m.jnt_pos[worldid, jntadr], xquat) + xpos xaxis = math.rot_vec_quat(jnt_axis, xquat) if jnt_type == wp.static(JointType.BALL.value): @@ -93,15 +93,15 @@ def _level(m: Model, d: Data, leveladr: int): ) xquat = math.mul_quat(xquat, qloc) # correct for off-center rotation - xpos = xanchor - math.rot_vec_quat(m.jnt_pos[jntadr], xquat) + xpos = xanchor - math.rot_vec_quat(m.jnt_pos[worldid, jntadr], xquat) elif jnt_type == wp.static(JointType.SLIDE.value): - xpos += xaxis * (qpos[qadr] - m.qpos0[qadr]) + xpos += xaxis * (qpos[qadr] - m.qpos0[worldid, qadr]) elif jnt_type == wp.static(JointType.HINGE.value): - qpos0 = m.qpos0[qadr] + qpos0 = m.qpos0[worldid, qadr] qloc = math.axis_angle_to_quat(jnt_axis, qpos[qadr] - qpos0) xquat = math.mul_quat(xquat, qloc) # correct for off-center rotation - xpos = xanchor - math.rot_vec_quat(m.jnt_pos[jntadr], xquat) + xpos = xanchor - math.rot_vec_quat(m.jnt_pos[worldid, jntadr], xquat) d.xanchor[worldid, jntadr] = xanchor d.xaxis[worldid, jntadr] = xaxis @@ -111,9 +111,11 @@ def _level(m: Model, d: Data, leveladr: int): xquat = wp.normalize(xquat) d.xquat[worldid, bodyid] = xquat d.xmat[worldid, bodyid] = math.quat_to_mat(xquat) - d.xipos[worldid, bodyid] = xpos + math.rot_vec_quat(m.body_ipos[bodyid], xquat) + d.xipos[worldid, bodyid] = xpos + math.rot_vec_quat( + m.body_ipos[worldid, bodyid], xquat + ) d.ximat[worldid, bodyid] = math.quat_to_mat( - math.mul_quat(xquat, m.body_iquat[bodyid]) + math.mul_quat(xquat, m.body_iquat[worldid, bodyid]) ) @kernel @@ -122,9 +124,11 @@ def geom_local_to_global(m: Model, d: Data): bodyid = m.geom_bodyid[geomid] xpos = d.xpos[worldid, bodyid] xquat = d.xquat[worldid, bodyid] - d.geom_xpos[worldid, geomid] = xpos + math.rot_vec_quat(m.geom_pos[geomid], xquat) + d.geom_xpos[worldid, geomid] = xpos + math.rot_vec_quat( + m.geom_pos[worldid, geomid], xquat + ) d.geom_xmat[worldid, geomid] = math.quat_to_mat( - math.mul_quat(xquat, m.geom_quat[geomid]) + math.mul_quat(xquat, m.geom_quat[worldid, geomid]) ) @kernel @@ -133,9 +137,11 @@ def site_local_to_global(m: Model, d: Data): bodyid = m.site_bodyid[siteid] xpos = d.xpos[worldid, bodyid] xquat = d.xquat[worldid, bodyid] - d.site_xpos[worldid, siteid] = xpos + math.rot_vec_quat(m.site_pos[siteid], xquat) + d.site_xpos[worldid, siteid] = xpos + math.rot_vec_quat( + m.site_pos[worldid, siteid], xquat + ) d.site_xmat[worldid, siteid] = math.quat_to_mat( - math.mul_quat(xquat, m.site_quat[siteid]) + math.mul_quat(xquat, m.site_quat[worldid, siteid]) ) wp.launch(_root, dim=(d.nworld), inputs=[m, d]) @@ -160,7 +166,9 @@ def com_pos(m: Model, d: Data): @kernel def subtree_com_init(m: Model, d: Data): worldid, bodyid = wp.tid() - d.subtree_com[worldid, bodyid] = d.xipos[worldid, bodyid] * m.body_mass[bodyid] + d.subtree_com[worldid, bodyid] = ( + d.xipos[worldid, bodyid] * m.body_mass[worldid, bodyid] + ) @kernel def subtree_com_acc(m: Model, d: Data, leveladr: int): @@ -172,14 +180,14 @@ def subtree_com_acc(m: Model, d: Data, leveladr: int): @kernel def subtree_div(m: Model, d: Data): worldid, bodyid = wp.tid() - d.subtree_com[worldid, bodyid] /= m.subtree_mass[bodyid] + d.subtree_com[worldid, bodyid] /= m.subtree_mass[worldid, bodyid] @kernel def cinert(m: Model, d: Data): worldid, bodyid = wp.tid() mat = d.ximat[worldid, bodyid] - inert = m.body_inertia[bodyid] - mass = m.body_mass[bodyid] + inert = m.body_inertia[worldid, bodyid] + mass = m.body_mass[worldid, bodyid] dif = d.xipos[worldid, bodyid] - d.subtree_com[worldid, m.body_rootid[bodyid]] # express inertia in com-based frame (mju_inertCom) @@ -264,9 +272,11 @@ def cam_local_to_global(m: Model, d: Data): bodyid = m.cam_bodyid[camid] xpos = d.xpos[worldid, bodyid] xquat = d.xquat[worldid, bodyid] - d.cam_xpos[worldid, camid] = xpos + math.rot_vec_quat(m.cam_pos[camid], xquat) + d.cam_xpos[worldid, camid] = xpos + math.rot_vec_quat( + m.cam_pos[worldid, camid], xquat + ) d.cam_xmat[worldid, camid] = math.quat_to_mat( - math.mul_quat(xquat, m.cam_quat[camid]) + math.mul_quat(xquat, m.cam_quat[worldid, camid]) ) @kernel @@ -280,10 +290,10 @@ def cam_fn(m: Model, d: Data): return elif m.cam_mode[camid] == wp.static(CamLightType.TRACK.value): body_xpos = d.xpos[worldid, m.cam_bodyid[camid]] - d.cam_xpos[worldid, camid] = body_xpos + m.cam_pos0[camid] + d.cam_xpos[worldid, camid] = body_xpos + m.cam_pos0[worldid, camid] elif m.cam_mode[camid] == wp.static(CamLightType.TRACKCOM.value): d.cam_xpos[worldid, camid] = ( - d.subtree_com[worldid, m.cam_bodyid[camid]] + m.cam_poscom0[camid] + d.subtree_com[worldid, m.cam_bodyid[camid]] + m.cam_poscom0[worldid, camid] ) elif m.cam_mode[camid] == wp.static(CamLightType.TARGETBODY.value) or m.cam_mode[ camid @@ -312,9 +322,11 @@ def light_local_to_global(m: Model, d: Data): xpos = d.xpos[worldid, bodyid] xquat = d.xquat[worldid, bodyid] d.light_xpos[worldid, lightid] = xpos + math.rot_vec_quat( - m.light_pos[lightid], xquat + m.light_pos[worldid, lightid], xquat + ) + d.light_xdir[worldid, lightid] = math.rot_vec_quat( + m.light_dir[worldid, lightid], xquat ) - d.light_xdir[worldid, lightid] = math.rot_vec_quat(m.light_dir[lightid], xquat) @kernel def light_fn(m: Model, d: Data): @@ -327,10 +339,11 @@ def light_fn(m: Model, d: Data): return elif m.light_mode[lightid] == wp.static(CamLightType.TRACK.value): body_xpos = d.xpos[worldid, m.light_bodyid[lightid]] - d.light_xpos[worldid, lightid] = body_xpos + m.light_pos0[lightid] + d.light_xpos[worldid, lightid] = body_xpos + m.light_pos0[worldid, lightid] elif m.light_mode[lightid] == wp.static(CamLightType.TRACKCOM.value): d.light_xpos[worldid, lightid] = ( - d.subtree_com[worldid, m.light_bodyid[lightid]] + m.light_poscom0[lightid] + d.subtree_com[worldid, m.light_bodyid[lightid]] + + m.light_poscom0[worldid, lightid] ) elif m.light_mode[lightid] == wp.static( CamLightType.TARGETBODY.value @@ -371,7 +384,7 @@ def qM_sparse(m: Model, d: Data): bodyid = m.dof_bodyid[dofid] # init M(i,i) with armature inertia - d.qM[worldid, 0, madr_ij] = m.dof_armature[dofid] + d.qM[worldid, 0, madr_ij] = m.dof_armature[worldid, dofid] # precompute buf = crb_body_i * cdof_i buf = math.inert_vec(d.crb[worldid, bodyid], d.cdof[worldid, dofid]) @@ -388,7 +401,7 @@ def qM_dense(m: Model, d: Data): bodyid = m.dof_bodyid[dofid] # init M(i,i) with armature inertia - M = m.dof_armature[dofid] + M = m.dof_armature[worldid, dofid] # precompute buf = crb_body_i * cdof_i buf = math.inert_vec(d.crb[worldid, bodyid], d.cdof[worldid, dofid]) @@ -689,7 +702,7 @@ def _cfrc_ext_equality(m: Model, d: Data): worldid = d.efc.worldid[efcid] id = d.efc.id[efcid] - eq_data = m.eq_data[id] + eq_data = m.eq_data[worldid, id] body_semantic = m.eq_objtype[id] == wp.static(ObjType.BODY.value) obj1 = m.eq_obj1id[id] @@ -710,7 +723,7 @@ def _cfrc_ext_equality(m: Model, d: Data): else: offset = wp.vec3(eq_data[3], eq_data[4], eq_data[5]) else: - offset = m.site_pos[obj1] + offset = m.site_pos[worldid, obj1] # transform point on body1: local -> global pos = d.xmat[worldid, bodyid1] @ offset + d.xpos[worldid, bodyid1] @@ -732,7 +745,7 @@ def _cfrc_ext_equality(m: Model, d: Data): else: offset = wp.vec3(eq_data[0], eq_data[1], eq_data[2]) else: - offset = m.site_pos[obj2] + offset = m.site_pos[worldid, obj2] # transform point on body2: local -> global pos = d.xmat[worldid, bodyid2] @ offset + d.xpos[worldid, bodyid2] @@ -811,7 +824,7 @@ def _transmission( ): worldid, actid = wp.tid() trntype = m.actuator_trntype[actid] - gear = m.actuator_gear[actid] + gear = m.actuator_gear[worldid, actid] if trntype == wp.static(TrnType.JOINT.value) or trntype == wp.static( TrnType.JOINTINPARENT.value ): @@ -1106,11 +1119,11 @@ def _forward(m: Model, d: Data): # update linear velocity lin -= wp.cross(xipos - subtree_com_root, ang) - d.subtree_linvel[worldid, bodyid] = m.body_mass[bodyid] * lin + d.subtree_linvel[worldid, bodyid] = m.body_mass[worldid, bodyid] * lin dv = wp.transpose(ximat) @ ang - dv[0] *= m.body_inertia[bodyid][0] - dv[1] *= m.body_inertia[bodyid][1] - dv[2] *= m.body_inertia[bodyid][2] + dv[0] *= m.body_inertia[worldid, bodyid][0] + dv[1] *= m.body_inertia[worldid, bodyid][1] + dv[2] *= m.body_inertia[worldid, bodyid][2] d.subtree_angmom[worldid, bodyid] = ximat @ dv d.subtree_bodyvel[worldid, bodyid] = wp.spatial_vector(ang, lin) @@ -1124,7 +1137,9 @@ def _linear_momentum(m: Model, d: Data, leveladr: int): if bodyid: pid = m.body_parentid[bodyid] wp.atomic_add(d.subtree_linvel[worldid], pid, d.subtree_linvel[worldid, bodyid]) - d.subtree_linvel[worldid, bodyid] /= wp.max(MJ_MINVAL, m.body_subtreemass[bodyid]) + d.subtree_linvel[worldid, bodyid] /= wp.max( + MJ_MINVAL, m.body_subtreemass[worldid, bodyid] + ) body_treeadr = m.body_treeadr.numpy() for i in reversed(range(len(body_treeadr))): @@ -1148,8 +1163,8 @@ def _angular_momentum(m: Model, d: Data, leveladr: int): vel = d.subtree_bodyvel[worldid, bodyid] linvel = d.subtree_linvel[worldid, bodyid] linvel_parent = d.subtree_linvel[worldid, pid] - mass = m.body_mass[bodyid] - subtreemass = m.body_subtreemass[bodyid] + mass = m.body_mass[worldid, bodyid] + subtreemass = m.body_subtreemass[worldid, bodyid] # momentum wrt body i dx = xipos - com @@ -1197,7 +1212,7 @@ def _joint_tendon(m: Model, d: Data): wrap_jnt_adr = m.wrap_jnt_adr[wrapid] wrap_objid = m.wrap_objid[wrap_jnt_adr] - prm = m.wrap_prm[wrap_jnt_adr] + prm = m.wrap_prm[worldid, wrap_jnt_adr] # add to length L = prm * d.qpos[worldid, m.jnt_qposadr[wrap_objid]] diff --git a/mujoco_warp/_src/support.py b/mujoco_warp/_src/support.py index 847bce8b..1d8d7de7 100644 --- a/mujoco_warp/_src/support.py +++ b/mujoco_warp/_src/support.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== -from typing import Tuple +from typing import Any, Tuple import mujoco import warp as wp diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index fef5b6c7..6b3a441f 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -444,7 +444,7 @@ class Constraint: gauss: gauss Cost (nworld,) cost: constraint + Gauss cost (nworld,) prev_cost: cost from previous iter (nworld,) - solver_niter: number of solver iterations (nworld,) + solver_niter: number of solver iterations (nworld, active: active (quadratic) constraints (njmax,) gtol: linesearch termination tolerance (nworld,) mv: qM @ search (nworld, nv) @@ -564,8 +564,8 @@ class Model: npair: number of predefined geom pairs () opt: physics options stat: model statistics - qpos0: qpos values at default pose (nq,) - qpos_spring: reference pose for springs (nq,) + qpos0: qpos values at default pose (nmodel, nq) + qpos_spring: reference pose for springs (nmodel, nq,) body_tree: BFS ordering of body ids body_treeadr: starting index of each body tree level actuator_moment_offset_nv: tiling configuration @@ -598,15 +598,15 @@ class Model: body_dofadr: start addr of dofs; -1: no dofs (nbody,) body_geomnum: number of geoms (nbody,) body_geomadr: start addr of geoms; -1: no geoms (nbody,) - body_pos: position offset rel. to parent body (nbody, 3) - body_quat: orientation offset rel. to parent body (nbody, 4) - body_ipos: local position of center of mass (nbody, 3) - body_iquat: local orientation of inertia ellipsoid (nbody, 4) - body_mass: mass (nbody,) - body_subtreemass: mass of subtree starting at this body (nbody,) - subtree_mass: mass of subtree (nbody,) - body_inertia: diagonal inertia in ipos/iquat frame (nbody, 3) - body_invweight0: mean inv inert in qpos0 (trn, rot) (nbody, 2) + body_pos: position offset rel. to parent body (nmodel, nbody, 3) + body_quat: orientation offset rel. to parent body (nmodel, nbody, 4) + body_ipos: local position of center of mass (nmodel, nbody, 3) + body_iquat: local orientation of inertia ellipsoid (nmodel, nbody, 4) + body_mass: mass (nmodel, nbody,) + body_subtreemass: mass of subtree starting at this body (nmodel, nbody,) + subtree_mass: mass of subtree (nmodel, nbody,) + body_inertia: diagonal inertia in ipos/iquat frame (nmodel, nbody, 3) + body_invweight0: mean inv inert in qpos0 (trn, rot) (nmodel, nbody, 2) body_contype: OR over all geom contypes (nbody,) body_conaffinity: OR over all geom conaffinities (nbody,) jnt_type: type of joint (mjtJoint) (njnt,) @@ -615,63 +615,63 @@ class Model: jnt_bodyid: id of joint's body (njnt,) jnt_limited: does joint have limits (njnt,) jnt_actfrclimited: does joint have actuator force limits (njnt,) - jnt_solref: constraint solver reference: limit (njnt, mjNREF) - jnt_solimp: constraint solver impedance: limit (njnt, mjNIMP) - jnt_pos: local anchor position (njnt, 3) + jnt_solref: constraint solver reference: limit (nmodel, njnt, mjNREF) + jnt_solimp: constraint solver impedance: limit (nmodel, njnt, mjNIMP) + jnt_pos: local anchor position (nmodel, njnt, 3) jnt_axis: local joint axis (njnt, 3) - jnt_stiffness: stiffness coefficient (njnt,) - jnt_range: joint limits (njnt, 2) - jnt_actfrcrange: range of total actuator force (njnt, 2) - jnt_margin: min distance for limit detection (njnt,) + jnt_stiffness: stiffness coefficient (nmodel, njnt,) + jnt_range: joint limits (nmodel, njnt, 2) + jnt_actfrcrange: range of total actuator force (nmodel, njnt, 2) + jnt_margin: min distance for limit detection (nmodel, njnt,) jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr jnt_limited_ball_adr: limited/ball jntadr dof_bodyid: id of dof's body (nv,) dof_jntid: id of dof's joint (nv,) dof_parentid: id of dof's parent; -1: none (nv,) dof_Madr: dof address in M-diagonal (nv,) - dof_armature: dof armature inertia/mass (nv,) - dof_damping: damping coefficient (nv,) - dof_invweight0: diag. inverse inertia in qpos0 (nv,) - dof_frictionloss: dof friction loss (nv,) - dof_solimp: constraint solver impedance: frictionloss (nv, NIMP) - dof_solref: constraint solver reference: frictionloss (nv, NREF) + dof_armature: dof armature inertia/mass (nmodel, nv) + dof_damping: damping coefficient (nmodel, nv) + dof_invweight0: diag. inverse inertia in qpos0 (nmodel, nv) + dof_frictionloss: dof friction loss (nmodel, nv) + dof_solimp: constraint solver impedance: frictionloss (nmodel, nv, NIMP) + dof_solref: constraint solver reference: frictionloss (nmodel, nv, NREF) dof_tri_row: np.tril_indices (mjm.nv)[0] dof_tri_col: np.tril_indices (mjm.nv)[1] geom_type: geometric type (mjtGeom) (ngeom,) - geom_contype: geom contact type (ngeom,) - geom_conaffinity: geom contact affinity (ngeom,) + geom_contype: geom contact type (nmodel, ngeom) + geom_conaffinity: geom contact affinity (nmodel, ngeom,) geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) geom_bodyid: id of geom's body (ngeom,) - geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) - geom_priority: geom contact priority (ngeom,) - geom_solmix: mixing coef for solref/imp in geom pair (ngeom,) - geom_solref: constraint solver reference: contact (ngeom, mjNREF) - geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) - geom_size: geom-specific size parameters (ngeom, 3) + geom_dataid: id of geom's mesh/hfield; -1: none (nmodel, ngeom) + geom_priority: geom contact priority (nmodel, ngeom) + geom_solmix: mixing coef for solref/imp in geom pair (nmodel, ngeom,) + geom_solref: constraint solver reference: contact (nmodel, ngeom, mjNREF) + geom_solimp: constraint solver impedance: contact (nmodel, ngeom, mjNIMP) + geom_size: geom-specific size parameters (nmodel, ngeom, 3) geom_aabb: bounding box, (center, size) (ngeom, 6) - geom_rbound: radius of bounding sphere (ngeom,) - geom_pos: local position offset rel. to body (ngeom, 3) - geom_quat: local orientation offset rel. to body (ngeom, 4) - geom_friction: friction for (slide, spin, roll) (ngeom, 3) - geom_margin: detect contact if dist