diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 77215478..66b6d7ac 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -233,8 +233,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data): 'collision_pair': d._impl.collision_pair.shape, 'collision_pairid': d._impl.collision_pairid.shape, 'collision_worldid': d._impl.collision_worldid.shape, - 'geom_xmat': d.geom_xmat.shape, - 'geom_xpos': d.geom_xpos.shape, 'nacon': d._impl.nacon.shape, 'ncollision': d._impl.ncollision.shape, 'contact__dim': d._impl.contact__dim.shape, @@ -253,15 +251,13 @@ def _collision_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _collision_shim, - num_outputs=20, + num_outputs=18, output_dims=output_dims, vmap_method=None, in_out_argnames={ 'collision_pair', 'collision_pairid', 'collision_worldid', - 'geom_xmat', - 'geom_xpos', 'nacon', 'ncollision', 'contact__dim', @@ -278,6 +274,27 @@ def _collision_jax_impl(m: types.Model, d: types.Data): 'contact__type', 'contact__worldid', }, + stage_in_argnames={ + 'geom_aabb', + 'geom_friction', + 'geom_gap', + 'geom_margin', + 'geom_rbound', + 'geom_size', + 'geom_solimp', + 'geom_solmix', + 'geom_solref', + 'geom_xmat', + 'geom_xpos', + 'hfield_data', + 'pair_friction', + 'pair_gap', + 'pair_margin', + 'pair_solimp', + 'pair_solref', + 'pair_solreffriction', + }, + stage_out_argnames={}, graph_mode=m.opt._impl.graph_mode, ) out = jf( @@ -373,23 +390,21 @@ def _collision_jax_impl(m: types.Model, d: types.Data): '_impl.collision_pair': out[0], '_impl.collision_pairid': out[1], '_impl.collision_worldid': out[2], - 'geom_xmat': out[3], - 'geom_xpos': out[4], - '_impl.nacon': out[5], - '_impl.ncollision': out[6], - '_impl.contact__dim': out[7], - '_impl.contact__dist': out[8], - '_impl.contact__frame': out[9], - '_impl.contact__friction': out[10], - '_impl.contact__geom': out[11], - '_impl.contact__geomcollisionid': out[12], - '_impl.contact__includemargin': out[13], - '_impl.contact__pos': out[14], - '_impl.contact__solimp': out[15], - '_impl.contact__solref': out[16], - '_impl.contact__solreffriction': out[17], - '_impl.contact__type': out[18], - '_impl.contact__worldid': out[19], + '_impl.nacon': out[3], + '_impl.ncollision': out[4], + '_impl.contact__dim': out[5], + '_impl.contact__dist': out[6], + '_impl.contact__frame': out[7], + '_impl.contact__friction': out[8], + '_impl.contact__geom': out[9], + '_impl.contact__geomcollisionid': out[10], + '_impl.contact__includemargin': out[11], + '_impl.contact__pos': out[12], + '_impl.contact__solimp': out[13], + '_impl.contact__solref': out[14], + '_impl.contact__solreffriction': out[15], + '_impl.contact__type': out[16], + '_impl.contact__worldid': out[17], }) return d diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index 0e92a3c8..fde118d3 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/mjx/mujoco/mjx/warp/ffi.py @@ -101,6 +101,8 @@ def jax_callable_variadic_tuple( vmap_method: Optional[str] = None, output_dims: Optional[dict[str, tuple[int, ...]]] = None, in_out_argnames: Optional[Sequence[str]] = None, + stage_in_argnames: Optional[Sequence[str]] = None, + stage_out_argnames: Optional[Sequence[str]] = None, ): """Wraps a JAX callable to support variadic tuples and dataclasses.""" @@ -130,6 +132,8 @@ def jax_callable_variadic_tuple( vmap_method=vmap_method, output_dims=output_dims, in_out_argnames=in_out_argnames, + stage_in_argnames=stage_in_argnames, + stage_out_argnames=stage_out_argnames, ) flat_args, in_tree = jax.tree.flatten(args) diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 74903557..4ccad4f5 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -952,7 +952,6 @@ def _forward_shim( def _forward_jax_impl(m: types.Model, d: types.Data): output_dims = { - 'act': d.act.shape, 'act_dot': d.act_dot.shape, 'actuator_force': d.actuator_force.shape, 'actuator_length': d.actuator_length.shape, @@ -970,10 +969,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'collision_pairid': d._impl.collision_pairid.shape, 'collision_worldid': d._impl.collision_worldid.shape, 'crb': d._impl.crb.shape, - 'ctrl': d.ctrl.shape, 'cvel': d.cvel.shape, 'energy': d._impl.energy.shape, - 'eq_active': d.eq_active.shape, 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, @@ -982,8 +979,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'geom_xpos': d.geom_xpos.shape, 'light_xdir': d._impl.light_xdir.shape, 'light_xpos': d._impl.light_xpos.shape, - 'mocap_pos': d.mocap_pos.shape, - 'mocap_quat': d.mocap_quat.shape, 'nacon': d._impl.nacon.shape, 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, @@ -1001,9 +996,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'qM': d._impl.qM.shape, 'qacc': d.qacc.shape, 'qacc_smooth': d.qacc_smooth.shape, - 'qacc_warmstart': d.qacc_warmstart.shape, 'qfrc_actuator': d.qfrc_actuator.shape, - 'qfrc_applied': d.qfrc_applied.shape, 'qfrc_bias': d.qfrc_bias.shape, 'qfrc_constraint': d.qfrc_constraint.shape, 'qfrc_damper': d._impl.qfrc_damper.shape, @@ -1012,7 +1005,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'qfrc_passive': d.qfrc_passive.shape, 'qfrc_smooth': d.qfrc_smooth.shape, 'qfrc_spring': d._impl.qfrc_spring.shape, - 'qpos': d.qpos.shape, 'qvel': d.qvel.shape, 'sensordata': d.sensordata.shape, 'site_xmat': d.site_xmat.shape, @@ -1027,12 +1019,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'ten_velocity': d._impl.ten_velocity.shape, 'ten_wrapadr': d._impl.ten_wrapadr.shape, 'ten_wrapnum': d._impl.ten_wrapnum.shape, - 'time': d.time.shape, 'wrap_obj': d._impl.wrap_obj.shape, 'wrap_xpos': d._impl.wrap_xpos.shape, 'xanchor': d.xanchor.shape, 'xaxis': d.xaxis.shape, - 'xfrc_applied': d.xfrc_applied.shape, 'ximat': d.ximat.shape, 'xipos': d.xipos.shape, 'xmat': d.xmat.shape, @@ -1085,11 +1075,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _forward_shim, - num_outputs=130, + num_outputs=120, output_dims=output_dims, vmap_method=None, in_out_argnames={ - 'act', 'act_dot', 'actuator_force', 'actuator_length', @@ -1107,10 +1096,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'collision_pairid', 'collision_worldid', 'crb', - 'ctrl', 'cvel', 'energy', - 'eq_active', 'flexedge_J', 'flexedge_length', 'flexedge_velocity', @@ -1119,8 +1106,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'geom_xpos', 'light_xdir', 'light_xpos', - 'mocap_pos', - 'mocap_quat', 'nacon', 'ncollision', 'ne', @@ -1138,9 +1123,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'qM', 'qacc', 'qacc_smooth', - 'qacc_warmstart', 'qfrc_actuator', - 'qfrc_applied', 'qfrc_bias', 'qfrc_constraint', 'qfrc_damper', @@ -1149,7 +1132,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'qfrc_passive', 'qfrc_smooth', 'qfrc_spring', - 'qpos', 'qvel', 'sensordata', 'site_xmat', @@ -1164,12 +1146,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'ten_velocity', 'ten_wrapadr', 'ten_wrapnum', - 'time', 'wrap_obj', 'wrap_xpos', 'xanchor', 'xaxis', - 'xfrc_applied', 'ximat', 'xipos', 'xmat', @@ -1220,6 +1200,175 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'efc__type', 'efc__vel', }, + stage_in_argnames={ + 'act', + 'act_dot', + 'actuator_acc0', + 'actuator_actrange', + 'actuator_biasprm', + 'actuator_ctrlrange', + 'actuator_dynprm', + 'actuator_force', + 'actuator_forcerange', + 'actuator_gainprm', + 'actuator_gear', + 'actuator_length', + 'body_gravcomp', + 'body_inertia', + 'body_invweight0', + 'body_ipos', + 'body_iquat', + 'body_mass', + 'body_pos', + 'body_quat', + 'body_subtreemass', + 'cam_mat0', + 'cam_pos', + 'cam_pos0', + 'cam_poscom0', + 'cam_quat', + 'cam_xmat', + 'cam_xpos', + 'cdof', + 'cdof_dot', + 'ctrl', + 'cvel', + 'dof_armature', + 'dof_damping', + 'dof_frictionloss', + 'dof_invweight0', + 'dof_solimp', + 'dof_solref', + 'eq_active', + 'eq_data', + 'eq_solimp', + 'eq_solref', + 'geom_aabb', + 'geom_friction', + 'geom_gap', + 'geom_margin', + 'geom_matid', + 'geom_pos', + 'geom_quat', + 'geom_rbound', + 'geom_rgba', + 'geom_size', + 'geom_solimp', + 'geom_solmix', + 'geom_solref', + 'geom_xmat', + 'geom_xpos', + 'hfield_data', + 'jnt_actfrcrange', + 'jnt_axis', + 'jnt_margin', + 'jnt_pos', + 'jnt_range', + 'jnt_solimp', + 'jnt_solref', + 'jnt_stiffness', + 'light_dir', + 'light_dir0', + 'light_pos', + 'light_pos0', + 'light_poscom0', + 'mat_rgba', + 'mocap_pos', + 'mocap_quat', + 'opt__density', + 'opt__gravity', + 'opt__ls_tolerance', + 'opt__magnetic', + 'opt__timestep', + 'opt__tolerance', + 'opt__viscosity', + 'opt__wind', + 'pair_friction', + 'pair_gap', + 'pair_margin', + 'pair_solimp', + 'pair_solref', + 'pair_solreffriction', + 'qacc', + 'qacc_smooth', + 'qacc_warmstart', + 'qfrc_actuator', + 'qfrc_applied', + 'qfrc_bias', + 'qfrc_constraint', + 'qfrc_fluid', + 'qfrc_gravcomp', + 'qfrc_passive', + 'qfrc_smooth', + 'qpos', + 'qpos0', + 'qpos_spring', + 'qvel', + 'sensordata', + 'site_pos', + 'site_quat', + 'site_xmat', + 'site_xpos', + 'subtree_com', + 'ten_length', + 'tendon_actfrcrange', + 'tendon_armature', + 'tendon_damping', + 'tendon_frictionloss', + 'tendon_invweight0', + 'tendon_length0', + 'tendon_lengthspring', + 'tendon_margin', + 'tendon_range', + 'tendon_solimp_fri', + 'tendon_solimp_lim', + 'tendon_solref_fri', + 'tendon_solref_lim', + 'tendon_stiffness', + 'time', + 'xanchor', + 'xaxis', + 'xfrc_applied', + 'ximat', + 'xipos', + 'xmat', + 'xpos', + 'xquat', + }, + stage_out_argnames={ + 'act_dot', + 'actuator_force', + 'actuator_length', + 'cam_xmat', + 'cam_xpos', + 'cdof', + 'cdof_dot', + 'cvel', + 'geom_xmat', + 'geom_xpos', + 'qacc', + 'qacc_smooth', + 'qfrc_actuator', + 'qfrc_bias', + 'qfrc_constraint', + 'qfrc_fluid', + 'qfrc_gravcomp', + 'qfrc_passive', + 'qfrc_smooth', + 'qvel', + 'sensordata', + 'site_xmat', + 'site_xpos', + 'subtree_com', + 'ten_length', + 'xanchor', + 'xaxis', + 'ximat', + 'xipos', + 'xmat', + 'xpos', + 'xquat', + }, graph_mode=m.opt._impl.graph_mode, ) out = jf( @@ -1673,136 +1822,126 @@ def _forward_jax_impl(m: types.Model, d: types.Data): d._impl.efc__vel, ) d = d.tree_replace({ - 'act': out[0], - 'act_dot': out[1], - 'actuator_force': out[2], - 'actuator_length': out[3], - '_impl.actuator_moment': out[4], - '_impl.actuator_velocity': out[5], - '_impl.cacc': out[6], - 'cam_xmat': out[7], - 'cam_xpos': out[8], - 'cdof': out[9], - 'cdof_dot': out[10], - '_impl.cfrc_ext': out[11], - '_impl.cfrc_int': out[12], - '_impl.cinert': out[13], - '_impl.collision_pair': out[14], - '_impl.collision_pairid': out[15], - '_impl.collision_worldid': out[16], - '_impl.crb': out[17], - 'ctrl': out[18], - 'cvel': out[19], - '_impl.energy': out[20], - 'eq_active': out[21], - '_impl.flexedge_J': out[22], - '_impl.flexedge_length': out[23], - '_impl.flexedge_velocity': out[24], - '_impl.flexvert_xpos': out[25], - 'geom_xmat': out[26], - 'geom_xpos': out[27], - '_impl.light_xdir': out[28], - '_impl.light_xpos': out[29], - 'mocap_pos': out[30], - 'mocap_quat': out[31], - '_impl.nacon': out[32], - '_impl.ncollision': out[33], - '_impl.ne': out[34], - '_impl.ne_connect': out[35], - '_impl.ne_flex': out[36], - '_impl.ne_jnt': out[37], - '_impl.ne_ten': out[38], - '_impl.ne_weld': out[39], - '_impl.nefc': out[40], - '_impl.nf': out[41], - '_impl.nl': out[42], - '_impl.nsolving': out[43], - '_impl.qLD': out[44], - '_impl.qLDiagInv': out[45], - '_impl.qM': out[46], - 'qacc': out[47], - 'qacc_smooth': out[48], - 'qacc_warmstart': out[49], - 'qfrc_actuator': out[50], - 'qfrc_applied': out[51], - 'qfrc_bias': out[52], - 'qfrc_constraint': out[53], - '_impl.qfrc_damper': out[54], - 'qfrc_fluid': out[55], - 'qfrc_gravcomp': out[56], - 'qfrc_passive': out[57], - 'qfrc_smooth': out[58], - '_impl.qfrc_spring': out[59], - 'qpos': out[60], - 'qvel': out[61], - 'sensordata': out[62], - 'site_xmat': out[63], - 'site_xpos': out[64], - '_impl.solver_niter': out[65], - '_impl.subtree_angmom': out[66], - '_impl.subtree_bodyvel': out[67], - 'subtree_com': out[68], - '_impl.subtree_linvel': out[69], - '_impl.ten_J': out[70], - 'ten_length': out[71], - '_impl.ten_velocity': out[72], - '_impl.ten_wrapadr': out[73], - '_impl.ten_wrapnum': out[74], - 'time': out[75], - '_impl.wrap_obj': out[76], - '_impl.wrap_xpos': out[77], - 'xanchor': out[78], - 'xaxis': out[79], - 'xfrc_applied': out[80], - 'ximat': out[81], - 'xipos': out[82], - 'xmat': out[83], - 'xpos': out[84], - 'xquat': out[85], - '_impl.contact__dim': out[86], - '_impl.contact__dist': out[87], - '_impl.contact__efc_address': out[88], - '_impl.contact__frame': out[89], - '_impl.contact__friction': out[90], - '_impl.contact__geom': out[91], - '_impl.contact__geomcollisionid': out[92], - '_impl.contact__includemargin': out[93], - '_impl.contact__pos': out[94], - '_impl.contact__solimp': out[95], - '_impl.contact__solref': out[96], - '_impl.contact__solreffriction': out[97], - '_impl.contact__type': out[98], - '_impl.contact__worldid': out[99], - '_impl.efc__D': out[100], - '_impl.efc__J': out[101], - '_impl.efc__Jaref': out[102], - '_impl.efc__Ma': out[103], - '_impl.efc__Mgrad': out[104], - '_impl.efc__alpha': out[105], - '_impl.efc__aref': out[106], - '_impl.efc__beta': out[107], - '_impl.efc__cost': out[108], - '_impl.efc__done': out[109], - '_impl.efc__force': out[110], - '_impl.efc__frictionloss': out[111], - '_impl.efc__gauss': out[112], - '_impl.efc__grad': out[113], - '_impl.efc__grad_dot': out[114], - '_impl.efc__id': out[115], - '_impl.efc__jv': out[116], - '_impl.efc__margin': out[117], - '_impl.efc__mv': out[118], - '_impl.efc__pos': out[119], - '_impl.efc__prev_Mgrad': out[120], - '_impl.efc__prev_cost': out[121], - '_impl.efc__prev_grad': out[122], - '_impl.efc__quad': out[123], - '_impl.efc__quad_gauss': out[124], - '_impl.efc__search': out[125], - '_impl.efc__search_dot': out[126], - '_impl.efc__state': out[127], - '_impl.efc__type': out[128], - '_impl.efc__vel': out[129], + 'act_dot': out[0], + 'actuator_force': out[1], + 'actuator_length': out[2], + '_impl.actuator_moment': out[3], + '_impl.actuator_velocity': out[4], + '_impl.cacc': out[5], + 'cam_xmat': out[6], + 'cam_xpos': out[7], + 'cdof': out[8], + 'cdof_dot': out[9], + '_impl.cfrc_ext': out[10], + '_impl.cfrc_int': out[11], + '_impl.cinert': out[12], + '_impl.collision_pair': out[13], + '_impl.collision_pairid': out[14], + '_impl.collision_worldid': out[15], + '_impl.crb': out[16], + 'cvel': out[17], + '_impl.energy': out[18], + '_impl.flexedge_J': out[19], + '_impl.flexedge_length': out[20], + '_impl.flexedge_velocity': out[21], + '_impl.flexvert_xpos': out[22], + 'geom_xmat': out[23], + 'geom_xpos': out[24], + '_impl.light_xdir': out[25], + '_impl.light_xpos': out[26], + '_impl.nacon': out[27], + '_impl.ncollision': out[28], + '_impl.ne': out[29], + '_impl.ne_connect': out[30], + '_impl.ne_flex': out[31], + '_impl.ne_jnt': out[32], + '_impl.ne_ten': out[33], + '_impl.ne_weld': out[34], + '_impl.nefc': out[35], + '_impl.nf': out[36], + '_impl.nl': out[37], + '_impl.nsolving': out[38], + '_impl.qLD': out[39], + '_impl.qLDiagInv': out[40], + '_impl.qM': out[41], + 'qacc': out[42], + 'qacc_smooth': out[43], + 'qfrc_actuator': out[44], + 'qfrc_bias': out[45], + 'qfrc_constraint': out[46], + '_impl.qfrc_damper': out[47], + 'qfrc_fluid': out[48], + 'qfrc_gravcomp': out[49], + 'qfrc_passive': out[50], + 'qfrc_smooth': out[51], + '_impl.qfrc_spring': out[52], + 'qvel': out[53], + 'sensordata': out[54], + 'site_xmat': out[55], + 'site_xpos': out[56], + '_impl.solver_niter': out[57], + '_impl.subtree_angmom': out[58], + '_impl.subtree_bodyvel': out[59], + 'subtree_com': out[60], + '_impl.subtree_linvel': out[61], + '_impl.ten_J': out[62], + 'ten_length': out[63], + '_impl.ten_velocity': out[64], + '_impl.ten_wrapadr': out[65], + '_impl.ten_wrapnum': out[66], + '_impl.wrap_obj': out[67], + '_impl.wrap_xpos': out[68], + 'xanchor': out[69], + 'xaxis': out[70], + 'ximat': out[71], + 'xipos': out[72], + 'xmat': out[73], + 'xpos': out[74], + 'xquat': out[75], + '_impl.contact__dim': out[76], + '_impl.contact__dist': out[77], + '_impl.contact__efc_address': out[78], + '_impl.contact__frame': out[79], + '_impl.contact__friction': out[80], + '_impl.contact__geom': out[81], + '_impl.contact__geomcollisionid': out[82], + '_impl.contact__includemargin': out[83], + '_impl.contact__pos': out[84], + '_impl.contact__solimp': out[85], + '_impl.contact__solref': out[86], + '_impl.contact__solreffriction': out[87], + '_impl.contact__type': out[88], + '_impl.contact__worldid': out[89], + '_impl.efc__D': out[90], + '_impl.efc__J': out[91], + '_impl.efc__Jaref': out[92], + '_impl.efc__Ma': out[93], + '_impl.efc__Mgrad': out[94], + '_impl.efc__alpha': out[95], + '_impl.efc__aref': out[96], + '_impl.efc__beta': out[97], + '_impl.efc__cost': out[98], + '_impl.efc__done': out[99], + '_impl.efc__force': out[100], + '_impl.efc__frictionloss': out[101], + '_impl.efc__gauss': out[102], + '_impl.efc__grad': out[103], + '_impl.efc__grad_dot': out[104], + '_impl.efc__id': out[105], + '_impl.efc__jv': out[106], + '_impl.efc__margin': out[107], + '_impl.efc__mv': out[108], + '_impl.efc__pos': out[109], + '_impl.efc__prev_Mgrad': out[110], + '_impl.efc__prev_cost': out[111], + '_impl.efc__prev_grad': out[112], + '_impl.efc__quad': out[113], + '_impl.efc__quad_gauss': out[114], + '_impl.efc__search': out[115], + '_impl.efc__search_dot': out[116], + '_impl.efc__state': out[117], + '_impl.efc__type': out[118], + '_impl.efc__vel': out[119], }) return d @@ -1818,25 +1957,6 @@ def forward_vmap(unused_axis_size, is_batched, m, d): return d, is_batched[1] -_m = mjwarp.Model( - **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} -) -_d = mjwarp.Data( - **{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init} -) -_o = mjwarp.Option( - **{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init} -) -_s = mjwarp.Statistic( - **{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init} -) -_c = mjwarp.Contact( - **{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init} -) -_e = mjwarp.Constraint( - **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} -) - @ffi.format_args_for_warp def _step_shim( # Model @@ -2769,10 +2889,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'collision_pairid': d._impl.collision_pairid.shape, 'collision_worldid': d._impl.collision_worldid.shape, 'crb': d._impl.crb.shape, - 'ctrl': d.ctrl.shape, 'cvel': d.cvel.shape, 'energy': d._impl.energy.shape, - 'eq_active': d.eq_active.shape, 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, @@ -2781,8 +2899,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'geom_xpos': d.geom_xpos.shape, 'light_xdir': d._impl.light_xdir.shape, 'light_xpos': d._impl.light_xpos.shape, - 'mocap_pos': d.mocap_pos.shape, - 'mocap_quat': d.mocap_quat.shape, 'nacon': d._impl.nacon.shape, 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, @@ -2802,7 +2918,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'qacc_smooth': d.qacc_smooth.shape, 'qacc_warmstart': d.qacc_warmstart.shape, 'qfrc_actuator': d.qfrc_actuator.shape, - 'qfrc_applied': d.qfrc_applied.shape, 'qfrc_bias': d.qfrc_bias.shape, 'qfrc_constraint': d.qfrc_constraint.shape, 'qfrc_damper': d._impl.qfrc_damper.shape, @@ -2831,7 +2946,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'wrap_xpos': d._impl.wrap_xpos.shape, 'xanchor': d.xanchor.shape, 'xaxis': d.xaxis.shape, - 'xfrc_applied': d.xfrc_applied.shape, 'ximat': d.ximat.shape, 'xipos': d.xipos.shape, 'xmat': d.xmat.shape, @@ -2884,7 +2998,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _step_shim, - num_outputs=130, + num_outputs=124, output_dims=output_dims, vmap_method=None, in_out_argnames={ @@ -2906,10 +3020,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'collision_pairid', 'collision_worldid', 'crb', - 'ctrl', 'cvel', 'energy', - 'eq_active', 'flexedge_J', 'flexedge_length', 'flexedge_velocity', @@ -2918,8 +3030,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'geom_xpos', 'light_xdir', 'light_xpos', - 'mocap_pos', - 'mocap_quat', 'nacon', 'ncollision', 'ne', @@ -2939,7 +3049,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'qacc_smooth', 'qacc_warmstart', 'qfrc_actuator', - 'qfrc_applied', 'qfrc_bias', 'qfrc_constraint', 'qfrc_damper', @@ -2968,7 +3077,6 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'wrap_xpos', 'xanchor', 'xaxis', - 'xfrc_applied', 'ximat', 'xipos', 'xmat', @@ -3019,6 +3127,179 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'efc__type', 'efc__vel', }, + stage_in_argnames={ + 'act', + 'act_dot', + 'actuator_acc0', + 'actuator_actrange', + 'actuator_biasprm', + 'actuator_ctrlrange', + 'actuator_dynprm', + 'actuator_force', + 'actuator_forcerange', + 'actuator_gainprm', + 'actuator_gear', + 'actuator_length', + 'body_gravcomp', + 'body_inertia', + 'body_invweight0', + 'body_ipos', + 'body_iquat', + 'body_mass', + 'body_pos', + 'body_quat', + 'body_subtreemass', + 'cam_mat0', + 'cam_pos', + 'cam_pos0', + 'cam_poscom0', + 'cam_quat', + 'cam_xmat', + 'cam_xpos', + 'cdof', + 'cdof_dot', + 'ctrl', + 'cvel', + 'dof_armature', + 'dof_damping', + 'dof_frictionloss', + 'dof_invweight0', + 'dof_solimp', + 'dof_solref', + 'eq_active', + 'eq_data', + 'eq_solimp', + 'eq_solref', + 'geom_aabb', + 'geom_friction', + 'geom_gap', + 'geom_margin', + 'geom_matid', + 'geom_pos', + 'geom_quat', + 'geom_rbound', + 'geom_rgba', + 'geom_size', + 'geom_solimp', + 'geom_solmix', + 'geom_solref', + 'geom_xmat', + 'geom_xpos', + 'hfield_data', + 'jnt_actfrcrange', + 'jnt_axis', + 'jnt_margin', + 'jnt_pos', + 'jnt_range', + 'jnt_solimp', + 'jnt_solref', + 'jnt_stiffness', + 'light_dir', + 'light_dir0', + 'light_pos', + 'light_pos0', + 'light_poscom0', + 'mat_rgba', + 'mocap_pos', + 'mocap_quat', + 'opt__density', + 'opt__gravity', + 'opt__ls_tolerance', + 'opt__magnetic', + 'opt__timestep', + 'opt__tolerance', + 'opt__viscosity', + 'opt__wind', + 'pair_friction', + 'pair_gap', + 'pair_margin', + 'pair_solimp', + 'pair_solref', + 'pair_solreffriction', + 'qacc', + 'qacc_smooth', + 'qacc_warmstart', + 'qfrc_actuator', + 'qfrc_applied', + 'qfrc_bias', + 'qfrc_constraint', + 'qfrc_fluid', + 'qfrc_gravcomp', + 'qfrc_passive', + 'qfrc_smooth', + 'qpos', + 'qpos0', + 'qpos_spring', + 'qvel', + 'sensordata', + 'site_pos', + 'site_quat', + 'site_xmat', + 'site_xpos', + 'subtree_com', + 'ten_length', + 'tendon_actfrcrange', + 'tendon_armature', + 'tendon_damping', + 'tendon_frictionloss', + 'tendon_invweight0', + 'tendon_length0', + 'tendon_lengthspring', + 'tendon_margin', + 'tendon_range', + 'tendon_solimp_fri', + 'tendon_solimp_lim', + 'tendon_solref_fri', + 'tendon_solref_lim', + 'tendon_stiffness', + 'time', + 'xanchor', + 'xaxis', + 'xfrc_applied', + 'ximat', + 'xipos', + 'xmat', + 'xpos', + 'xquat', + }, + stage_out_argnames={ + 'act', + 'act_dot', + 'actuator_force', + 'actuator_length', + 'cam_xmat', + 'cam_xpos', + 'cdof', + 'cdof_dot', + 'cvel', + 'geom_xmat', + 'geom_xpos', + 'qacc', + 'qacc_smooth', + 'qacc_warmstart', + 'qfrc_actuator', + 'qfrc_bias', + 'qfrc_constraint', + 'qfrc_fluid', + 'qfrc_gravcomp', + 'qfrc_passive', + 'qfrc_smooth', + 'qpos', + 'qvel', + 'sensordata', + 'site_xmat', + 'site_xpos', + 'subtree_com', + 'ten_length', + 'time', + 'xanchor', + 'xaxis', + 'ximat', + 'xipos', + 'xmat', + 'xpos', + 'xquat', + }, graph_mode=m.opt._impl.graph_mode, ) out = jf( @@ -3492,118 +3773,112 @@ def _step_jax_impl(m: types.Model, d: types.Data): '_impl.collision_pairid': out[15], '_impl.collision_worldid': out[16], '_impl.crb': out[17], - 'ctrl': out[18], - 'cvel': out[19], - '_impl.energy': out[20], - 'eq_active': out[21], - '_impl.flexedge_J': out[22], - '_impl.flexedge_length': out[23], - '_impl.flexedge_velocity': out[24], - '_impl.flexvert_xpos': out[25], - 'geom_xmat': out[26], - 'geom_xpos': out[27], - '_impl.light_xdir': out[28], - '_impl.light_xpos': out[29], - 'mocap_pos': out[30], - 'mocap_quat': out[31], - '_impl.nacon': out[32], - '_impl.ncollision': out[33], - '_impl.ne': out[34], - '_impl.ne_connect': out[35], - '_impl.ne_flex': out[36], - '_impl.ne_jnt': out[37], - '_impl.ne_ten': out[38], - '_impl.ne_weld': out[39], - '_impl.nefc': out[40], - '_impl.nf': out[41], - '_impl.nl': out[42], - '_impl.nsolving': out[43], - '_impl.qLD': out[44], - '_impl.qLDiagInv': out[45], - '_impl.qM': out[46], - 'qacc': out[47], - 'qacc_smooth': out[48], - 'qacc_warmstart': out[49], - 'qfrc_actuator': out[50], - 'qfrc_applied': out[51], - 'qfrc_bias': out[52], - 'qfrc_constraint': out[53], - '_impl.qfrc_damper': out[54], - 'qfrc_fluid': out[55], - 'qfrc_gravcomp': out[56], - 'qfrc_passive': out[57], - 'qfrc_smooth': out[58], - '_impl.qfrc_spring': out[59], - 'qpos': out[60], - 'qvel': out[61], - 'sensordata': out[62], - 'site_xmat': out[63], - 'site_xpos': out[64], - '_impl.solver_niter': out[65], - '_impl.subtree_angmom': out[66], - '_impl.subtree_bodyvel': out[67], - 'subtree_com': out[68], - '_impl.subtree_linvel': out[69], - '_impl.ten_J': out[70], - 'ten_length': out[71], - '_impl.ten_velocity': out[72], - '_impl.ten_wrapadr': out[73], - '_impl.ten_wrapnum': out[74], - 'time': out[75], - '_impl.wrap_obj': out[76], - '_impl.wrap_xpos': out[77], - 'xanchor': out[78], - 'xaxis': out[79], - 'xfrc_applied': out[80], - 'ximat': out[81], - 'xipos': out[82], - 'xmat': out[83], - 'xpos': out[84], - 'xquat': out[85], - '_impl.contact__dim': out[86], - '_impl.contact__dist': out[87], - '_impl.contact__efc_address': out[88], - '_impl.contact__frame': out[89], - '_impl.contact__friction': out[90], - '_impl.contact__geom': out[91], - '_impl.contact__geomcollisionid': out[92], - '_impl.contact__includemargin': out[93], - '_impl.contact__pos': out[94], - '_impl.contact__solimp': out[95], - '_impl.contact__solref': out[96], - '_impl.contact__solreffriction': out[97], - '_impl.contact__type': out[98], - '_impl.contact__worldid': out[99], - '_impl.efc__D': out[100], - '_impl.efc__J': out[101], - '_impl.efc__Jaref': out[102], - '_impl.efc__Ma': out[103], - '_impl.efc__Mgrad': out[104], - '_impl.efc__alpha': out[105], - '_impl.efc__aref': out[106], - '_impl.efc__beta': out[107], - '_impl.efc__cost': out[108], - '_impl.efc__done': out[109], - '_impl.efc__force': out[110], - '_impl.efc__frictionloss': out[111], - '_impl.efc__gauss': out[112], - '_impl.efc__grad': out[113], - '_impl.efc__grad_dot': out[114], - '_impl.efc__id': out[115], - '_impl.efc__jv': out[116], - '_impl.efc__margin': out[117], - '_impl.efc__mv': out[118], - '_impl.efc__pos': out[119], - '_impl.efc__prev_Mgrad': out[120], - '_impl.efc__prev_cost': out[121], - '_impl.efc__prev_grad': out[122], - '_impl.efc__quad': out[123], - '_impl.efc__quad_gauss': out[124], - '_impl.efc__search': out[125], - '_impl.efc__search_dot': out[126], - '_impl.efc__state': out[127], - '_impl.efc__type': out[128], - '_impl.efc__vel': out[129], + 'cvel': out[18], + '_impl.energy': out[19], + '_impl.flexedge_J': out[20], + '_impl.flexedge_length': out[21], + '_impl.flexedge_velocity': out[22], + '_impl.flexvert_xpos': out[23], + 'geom_xmat': out[24], + 'geom_xpos': out[25], + '_impl.light_xdir': out[26], + '_impl.light_xpos': out[27], + '_impl.nacon': out[28], + '_impl.ncollision': out[29], + '_impl.ne': out[30], + '_impl.ne_connect': out[31], + '_impl.ne_flex': out[32], + '_impl.ne_jnt': out[33], + '_impl.ne_ten': out[34], + '_impl.ne_weld': out[35], + '_impl.nefc': out[36], + '_impl.nf': out[37], + '_impl.nl': out[38], + '_impl.nsolving': out[39], + '_impl.qLD': out[40], + '_impl.qLDiagInv': out[41], + '_impl.qM': out[42], + 'qacc': out[43], + 'qacc_smooth': out[44], + 'qacc_warmstart': out[45], + 'qfrc_actuator': out[46], + 'qfrc_bias': out[47], + 'qfrc_constraint': out[48], + '_impl.qfrc_damper': out[49], + 'qfrc_fluid': out[50], + 'qfrc_gravcomp': out[51], + 'qfrc_passive': out[52], + 'qfrc_smooth': out[53], + '_impl.qfrc_spring': out[54], + 'qpos': out[55], + 'qvel': out[56], + 'sensordata': out[57], + 'site_xmat': out[58], + 'site_xpos': out[59], + '_impl.solver_niter': out[60], + '_impl.subtree_angmom': out[61], + '_impl.subtree_bodyvel': out[62], + 'subtree_com': out[63], + '_impl.subtree_linvel': out[64], + '_impl.ten_J': out[65], + 'ten_length': out[66], + '_impl.ten_velocity': out[67], + '_impl.ten_wrapadr': out[68], + '_impl.ten_wrapnum': out[69], + 'time': out[70], + '_impl.wrap_obj': out[71], + '_impl.wrap_xpos': out[72], + 'xanchor': out[73], + 'xaxis': out[74], + 'ximat': out[75], + 'xipos': out[76], + 'xmat': out[77], + 'xpos': out[78], + 'xquat': out[79], + '_impl.contact__dim': out[80], + '_impl.contact__dist': out[81], + '_impl.contact__efc_address': out[82], + '_impl.contact__frame': out[83], + '_impl.contact__friction': out[84], + '_impl.contact__geom': out[85], + '_impl.contact__geomcollisionid': out[86], + '_impl.contact__includemargin': out[87], + '_impl.contact__pos': out[88], + '_impl.contact__solimp': out[89], + '_impl.contact__solref': out[90], + '_impl.contact__solreffriction': out[91], + '_impl.contact__type': out[92], + '_impl.contact__worldid': out[93], + '_impl.efc__D': out[94], + '_impl.efc__J': out[95], + '_impl.efc__Jaref': out[96], + '_impl.efc__Ma': out[97], + '_impl.efc__Mgrad': out[98], + '_impl.efc__alpha': out[99], + '_impl.efc__aref': out[100], + '_impl.efc__beta': out[101], + '_impl.efc__cost': out[102], + '_impl.efc__done': out[103], + '_impl.efc__force': out[104], + '_impl.efc__frictionloss': out[105], + '_impl.efc__gauss': out[106], + '_impl.efc__grad': out[107], + '_impl.efc__grad_dot': out[108], + '_impl.efc__id': out[109], + '_impl.efc__jv': out[110], + '_impl.efc__margin': out[111], + '_impl.efc__mv': out[112], + '_impl.efc__pos': out[113], + '_impl.efc__prev_Mgrad': out[114], + '_impl.efc__prev_cost': out[115], + '_impl.efc__prev_grad': out[116], + '_impl.efc__quad': out[117], + '_impl.efc__quad_gauss': out[118], + '_impl.efc__search': out[119], + '_impl.efc__search_dot': out[120], + '_impl.efc__state': out[121], + '_impl.efc__type': out[122], + '_impl.efc__vel': out[123], }) return d diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 18fee61e..fe65576b 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -218,7 +218,7 @@ class StepTest(parameterized.TestCase): batch_size=(1, 7), # NOTE: GraphMode.JAX is incompatible with MuJoCo Warp at the moment, # even when setting graph_conditional=False. - graph_mode=('WARP',), + graph_mode=('WARP', 'WARP_STAGED'), ) def test_step(self, xml: str, batch_size: int, graph_mode: str): if not _FORCE_TEST: diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 7c018b49..ff12107c 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -136,9 +136,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): output_dims = { 'geom_xmat': d.geom_xmat.shape, 'geom_xpos': d.geom_xpos.shape, - 'mocap_pos': d.mocap_pos.shape, - 'mocap_quat': d.mocap_quat.shape, - 'qpos': d.qpos.shape, 'site_xmat': d.site_xmat.shape, 'site_xpos': d.site_xpos.shape, 'xanchor': d.xanchor.shape, @@ -151,15 +148,52 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _kinematics_shim, - num_outputs=14, + num_outputs=11, output_dims=output_dims, vmap_method=None, in_out_argnames={ 'geom_xmat', 'geom_xpos', + 'site_xmat', + 'site_xpos', + 'xanchor', + 'xaxis', + 'ximat', + 'xipos', + 'xmat', + 'xpos', + 'xquat', + }, + stage_in_argnames={ + 'body_ipos', + 'body_iquat', + 'body_pos', + 'body_quat', + 'geom_pos', + 'geom_quat', + 'geom_xmat', + 'geom_xpos', + 'jnt_axis', + 'jnt_pos', 'mocap_pos', 'mocap_quat', 'qpos', + 'qpos0', + 'site_pos', + 'site_quat', + 'site_xmat', + 'site_xpos', + 'xanchor', + 'xaxis', + 'ximat', + 'xipos', + 'xmat', + 'xpos', + 'xquat', + }, + stage_out_argnames={ + 'geom_xmat', + 'geom_xpos', 'site_xmat', 'site_xpos', 'xanchor', @@ -216,18 +250,15 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): d = d.tree_replace({ 'geom_xmat': out[0], 'geom_xpos': out[1], - 'mocap_pos': out[2], - 'mocap_quat': out[3], - 'qpos': out[4], - 'site_xmat': out[5], - 'site_xpos': out[6], - 'xanchor': out[7], - 'xaxis': out[8], - 'ximat': out[9], - 'xipos': out[10], - 'xmat': out[11], - 'xpos': out[12], - 'xquat': out[13], + 'site_xmat': out[2], + 'site_xpos': out[3], + 'xanchor': out[4], + 'xaxis': out[5], + 'ximat': out[6], + 'xipos': out[7], + 'xmat': out[8], + 'xpos': out[9], + 'xquat': out[10], }) return d @@ -243,25 +274,6 @@ def kinematics_vmap(unused_axis_size, is_batched, m, d): return d, is_batched[1] -_m = mjwarp.Model( - **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} -) -_d = mjwarp.Data( - **{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init} -) -_o = mjwarp.Option( - **{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init} -) -_s = mjwarp.Statistic( - **{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init} -) -_c = mjwarp.Contact( - **{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init} -) -_e = mjwarp.Constraint( - **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} -) - @ffi.format_args_for_warp def _tendon_shim( # Model @@ -348,12 +360,6 @@ def _tendon_shim( def _tendon_jax_impl(m: types.Model, d: types.Data): output_dims = { - 'cdof': d.cdof.shape, - 'geom_xmat': d.geom_xmat.shape, - 'geom_xpos': d.geom_xpos.shape, - 'qpos': d.qpos.shape, - 'site_xpos': d.site_xpos.shape, - 'subtree_com': d.subtree_com.shape, 'ten_J': d._impl.ten_J.shape, 'ten_length': d.ten_length.shape, 'ten_wrapadr': d._impl.ten_wrapadr.shape, @@ -363,16 +369,10 @@ def _tendon_jax_impl(m: types.Model, d: types.Data): } jf = ffi.jax_callable_variadic_tuple( _tendon_shim, - num_outputs=12, + num_outputs=6, output_dims=output_dims, vmap_method=None, in_out_argnames={ - 'cdof', - 'geom_xmat', - 'geom_xpos', - 'qpos', - 'site_xpos', - 'subtree_com', 'ten_J', 'ten_length', 'ten_wrapadr', @@ -380,6 +380,17 @@ def _tendon_jax_impl(m: types.Model, d: types.Data): 'wrap_obj', 'wrap_xpos', }, + stage_in_argnames={ + 'cdof', + 'geom_size', + 'geom_xmat', + 'geom_xpos', + 'qpos', + 'site_xpos', + 'subtree_com', + 'ten_length', + }, + stage_out_argnames={'ten_length'}, graph_mode=m.opt._impl.graph_mode, ) out = jf( @@ -421,18 +432,12 @@ def _tendon_jax_impl(m: types.Model, d: types.Data): d._impl.wrap_xpos, ) d = d.tree_replace({ - 'cdof': out[0], - 'geom_xmat': out[1], - 'geom_xpos': out[2], - 'qpos': out[3], - 'site_xpos': out[4], - 'subtree_com': out[5], - '_impl.ten_J': out[6], - 'ten_length': out[7], - '_impl.ten_wrapadr': out[8], - '_impl.ten_wrapnum': out[9], - '_impl.wrap_obj': out[10], - '_impl.wrap_xpos': out[11], + '_impl.ten_J': out[0], + 'ten_length': out[1], + '_impl.ten_wrapadr': out[2], + '_impl.ten_wrapnum': out[3], + '_impl.wrap_obj': out[4], + '_impl.wrap_xpos': out[5], }) return d diff --git a/mjx/mujoco/mjx/warp/testspeed.py b/mjx/mujoco/mjx/warp/testspeed.py index fef2a386..69d42517 100644 --- a/mjx/mujoco/mjx/warp/testspeed.py +++ b/mjx/mujoco/mjx/warp/testspeed.py @@ -22,6 +22,7 @@ from typing import Any, Callable, Sequence, Tuple from absl import app from absl import flags import jax +import jax.numpy as jnp import mujoco from mujoco import mjx from mujoco.mjx._src import test_util @@ -29,6 +30,7 @@ from mujoco.mjx.warp import collision_driver as wp_collision from mujoco.mjx.warp import forward as wp_forward from mujoco.mjx.warp import smooth as wp_smooth import mujoco.mjx.third_party.mujoco_warp as mjwarp +import numpy as np import warp as wp from mujoco.mjx.third_party.warp._src.jax_experimental import ffi as warp_ffi @@ -50,6 +52,18 @@ _WP_KERNEL_CACHE_DIR = flags.DEFINE_string( None, 'Path to the Warp kernel cache directory.', ) +_GRAPH_MODE = flags.DEFINE_enum( + 'graph_mode', + 'WARP', + ['NONE', 'WARP', 'WARP_STAGED', 'WARP_STAGED_EX'], + 'Graph capture mode for JAX WARP FFI benchmark.', +) +_BENCHMARK = flags.DEFINE_enum( + 'benchmark', + 'jax_warp', + ['jax_warp', 'jax', 'warp'], + 'Which benchmark to run.', +) _COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1} jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS) @@ -95,10 +109,14 @@ def benchmark( d = mjx.make_data( m, impl=mx.impl, naconmax=_NACONMAX.value, njmax=_NJMAX.value ) + # Initialize from first keyframe if available + if m.nkey > 0: + d = d.replace(qpos=m.key_qpos[0], ctrl=m.key_ctrl[0]) return d key = jax.random.split(jax.random.key(0), nenv) d = jax_jit(init)(key) + jax.block_until_ready(d) @jax_jit @@ -115,177 +133,6 @@ def benchmark( return jit_time, run_time, steps -def benchmark_raw_jax_warp( - m: mujoco.MjModel, - nstep: int = 1000, - nenv: int = 8192, - unroll_steps: int = 4, - function: str = 'kinematics', -): - if function not in ('kinematics', 'forward', 'step', 'collision'): - raise NotImplementedError( - f'{function} is not implemented for raw warp speed test.' - ) - - def warp_fn( - qpos_in: wp.array2d(dtype=wp.float32), - qpos_out: wp.array2d(dtype=wp.float32), - time_: wp.array(dtype=wp.float32), - xpos: wp.array2d(dtype=wp.vec3), - xquat: wp.array2d(dtype=wp.quat), - xmat: wp.array2d(dtype=wp.mat33), - xipos: wp.array2d(dtype=wp.vec3), - ximat: wp.array2d(dtype=wp.mat33), - xanchor: wp.array2d(dtype=wp.vec3), - xaxis: wp.array2d(dtype=wp.vec3), - geom_xpos: wp.array2d(dtype=wp.vec3), - geom_xmat: wp.array2d(dtype=wp.mat33), - ): - wp.copy(d.qpos, qpos_in) - if function == 'kinematics': - mjwarp.kinematics(mw, d) - elif function == 'forward': - mjwarp.forward(mw, d) - elif function == 'step': - mjwarp.step(mw, d) - elif function == 'collision': - mjwarp.collision(mw, d) - else: - raise NotImplementedError(f'{function} not implemented in speed test.') - wp.copy(qpos_out, d.qpos) - wp.copy(time_, d.time) - wp.copy(xpos, d.xpos) - wp.copy(xquat, d.xquat) - wp.copy(xmat, d.xmat) - wp.copy(xipos, d.xipos) - wp.copy(ximat, d.ximat) - wp.copy(xanchor, d.xanchor) - wp.copy(xaxis, d.xaxis) - wp.copy(geom_xpos, d.geom_xpos) - wp.copy(geom_xmat, d.geom_xmat) - - def unroll( - qpos, - qpos_out, - time_, - xpos, - xquat, - xmat, - xipos, - ximat, - xanchor, - xaxis, - geom_xpos, - geom_xmat, - ): - def step(carry, _): - qpos_, *_ = carry - out = warp_fn_jax(qpos_ + 0.0 * qpos_) - out = tuple(out) - return (out[0],) + out, None - - ( - qpos, - qpos_out, - time_, - xpos, - xquat, - xmat, - xipos, - ximat, - xanchor, - xaxis, - geom_xpos, - geom_xmat, - ), _ = jax.lax.scan( - step, - ( - qpos, - qpos_out, - time_, - xpos, - xquat, - xmat, - xipos, - ximat, - xanchor, - xaxis, - geom_xpos, - geom_xmat, - ), - length=nstep, - unroll=unroll_steps, - ) - - return ( - qpos, - qpos_out, - time_, - xpos, - xquat, - xmat, - xipos, - ximat, - xanchor, - xaxis, - geom_xpos, - geom_xmat, - ) - - output_dims = { - 'time_': (nenv,), - 'qpos_out': (nenv, m.nq), - 'xpos': (nenv, m.nbody, 3), - 'xquat': (nenv, m.nbody, 4), - 'xmat': (nenv, m.nbody, 3, 3), - 'xipos': (nenv, m.nbody, 3), - 'ximat': (nenv, m.nbody, 3, 3), - 'xanchor': (nenv, m.njnt, 3), - 'xaxis': (nenv, m.njnt, 3), - 'geom_xpos': (nenv, m.ngeom, 3), - 'geom_xmat': (nenv, m.ngeom, 3, 3), - } - warp_fn_jax = warp_ffi.jax_callable( - warp_fn, - num_outputs=11, - output_dims=output_dims, - graph_mode=warp_ffi.GraphMode.JAX, - ) - - @jax.vmap - def init(key): - d = mjx.make_data(m, impl='jax') - return d - - key = jax.random.split(jax.random.key(0), nenv) - dx = jax_jit(init)(key) - d_ = mujoco.MjData(m) - mw = mjwarp.put_model(m) - d = mjwarp.put_data( - m, d_, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value - ) - - jax_unroll_fn = jax_jit(unroll) - jit_time, run_time = _measure( - jax_unroll_fn, - dx.qpos, - dx.qpos, - dx.time, - dx.xpos, - dx.xquat, - dx.xmat, - dx.xipos, - dx.ximat, - dx.xanchor, - dx.xaxis, - dx.geom_xpos, - dx.geom_xmat, - ) - steps = nstep * nenv - - return jit_time, run_time, steps - - def _compile_fn(fn, m, d): fn(m, d) fn(m, d) @@ -310,9 +157,16 @@ def benchmark_raw_warp( mw = mjwarp.put_model(m) dw = mjwarp.make_data( - m, nworld=nenv, nconmax=_NACONMAX.value, njmax=_NJMAX.value + m, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value ) + # Initialize from first keyframe if available + if m.nkey > 0: + qpos_init = np.tile(m.key_qpos[0], (nenv, 1)).astype(np.float32) + ctrl_init = np.tile(m.key_ctrl[0], (nenv, 1)).astype(np.float32) + wp.copy(dw.qpos, wp.from_numpy(qpos_init)) + wp.copy(dw.ctrl, wp.from_numpy(ctrl_init)) + if function == 'kinematics': fn = mjwarp.kinematics elif function == 'forward': @@ -353,8 +207,16 @@ def _main(_: Sequence[str]): except Exception as _: m = mujoco.MjModel.from_xml_path(modelfile) - mx = mjx.put_model(m, impl='jax') - mw = mjx.put_model(m, impl='warp') + benchmark_type = _BENCHMARK.value + graph_mode = getattr(warp_ffi.GraphMode, _GRAPH_MODE.value) + + # Only allocate the model needed for the specific benchmark + mx = None + mw = None + if benchmark_type == 'jax_warp': + mw = mjx.put_model(m, impl='warp', graph_mode=graph_mode) + elif benchmark_type == 'jax': + mx = mjx.put_model(m, impl='jax') if function_ == 'kinematics': func_warp = jax.vmap(wp_smooth.kinematics, in_axes=(None, 0)) @@ -377,38 +239,24 @@ def _main(_: Sequence[str]): print(f' nenv : {nenv}') print(f' nstep : {nstep}') print(f' timestep : {m.opt.timestep}') - print(f' unroll : {unroll}\n') + print(f' unroll : {unroll}') + print(f' benchmark : {benchmark_type}') + print(f' graph_mode : {_GRAPH_MODE.value}\n') - for name, mx_, op in ( - ('JAX WARP FFI', mw, func_warp), - ('Pure JAX', mx, func_jax), - ): - if op is not None: - jit_time, run_time, steps = benchmark(m, mx_, op, nstep, nenv, unroll) + if benchmark_type == 'jax_warp': + jit_time, run_time, steps = benchmark(m, mw, func_warp, nstep, nenv, unroll) + print(f' JAX WARP FFI (GraphMode: {_GRAPH_MODE.value}):') + elif benchmark_type == 'jax': + jit_time, run_time, steps = benchmark(m, mx, func_jax, nstep, nenv, unroll) + print(' Pure JAX:') + elif benchmark_type == 'warp': + jit_time, run_time, steps = benchmark_raw_warp( + m, nstep, nenv, unroll, function=function_ + ) + print(' Pure WARP:') + else: + raise ValueError(f'Unknown benchmark type: {benchmark_type}') - print(f' {name}:') - print(f' JIT time : {jit_time:.2f} s') - print(f' simulation time : {run_time:.2f} s') - print(f' steps per second : {steps / run_time:,.0f}') - print( - f' realtime factor : {steps * m.opt.timestep / run_time:.2f} x' - ) - print(f' time per step : {1e6 * run_time / steps:.2f} µs\n') - - jit_time, run_time, steps = benchmark_raw_jax_warp( - m, nstep, nenv, unroll, function=function_ - ) - print(' Pure JAX-WARP:') - print(f' JIT time : {jit_time:.2f} s') - print(f' simulation time : {run_time:.2f} s') - print(f' steps per second : {steps / run_time:,.0f}') - print(f' realtime factor : {steps * m.opt.timestep / run_time:.2f} x') - print(f' time per step : {1e6 * run_time / steps:.2f} µs\n') - - jit_time, run_time, steps = benchmark_raw_warp( - m, nstep, nenv, unroll, function=function_ - ) - print(' Pure WARP:') print(f' JIT time : {jit_time:.2f} s') print(f' simulation time : {run_time:.2f} s') print(f' steps per second : {steps / run_time:,.0f}') diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index f323673e..29dc02c3 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -23,7 +23,6 @@ from jax import tree_util from jax.interpreters import batching from mujoco.mjx._src import dataclasses as mjx_dataclasses import numpy as np - if typing.TYPE_CHECKING: GraphMode = int else: @@ -44,6 +43,7 @@ class TileSet: adr: address of each tile in the set size: size of all the tiles in this set """ + adr: np.ndarray size: int