diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py index 09a0a512..a13fc55e 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py @@ -799,61 +799,39 @@ def _tendon_actuator_force_clamp( actuator_force_out[worldid, actid] *= actfrcrange[1] / ten_actfrc -def _qfrc_actuator(m: Model, d: Data): - NU = m.nu +@wp.kernel +def _qfrc_actuator( + # Model: + nu: int, + ngravcomp: int, + jnt_actfrclimited: wp.array(dtype=bool), + jnt_actfrcrange: wp.array2d(dtype=wp.vec2), + jnt_actgravcomp: wp.array(dtype=int), + dof_jntid: wp.array(dtype=int), + # Data in: + actuator_moment_in: wp.array3d(dtype=float), + qfrc_gravcomp_in: wp.array2d(dtype=float), + actuator_force_in: wp.array2d(dtype=float), + # Data out: + qfrc_actuator_out: wp.array2d(dtype=float), +): + worldid, dofid = wp.tid() - @wp.kernel - def qfrc_actuator( - # Model: - ngravcomp: int, - jnt_actfrclimited: wp.array(dtype=bool), - jnt_actfrcrange: wp.array2d(dtype=wp.vec2), - jnt_actgravcomp: wp.array(dtype=int), - dof_jntid: wp.array(dtype=int), - # Data in: - actuator_moment_in: wp.array3d(dtype=float), - qfrc_gravcomp_in: wp.array2d(dtype=float), - actuator_force_in: wp.array2d(dtype=float), - # Data out: - qfrc_actuator_out: wp.array2d(dtype=float), - ): - worldid, dofid = wp.tid() + qfrc = float(0.0) + for uid in range(nu): + qfrc += actuator_moment_in[worldid, uid, dofid] * actuator_force_in[worldid, uid] - actuator_moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, 1), offset=(0, dofid)) - actuator_moment_tile = wp.tile_squeeze(actuator_moment_tile, axis=(1,)) - actuator_force_tile = wp.tile_load(actuator_force_in[worldid], shape=NU) - actuator_moment_force_tile = wp.tile_map(wp.mul, actuator_moment_tile, actuator_force_tile) - qfrc_tile = wp.tile_reduce(wp.add, actuator_moment_force_tile) - qfrc = qfrc_tile[0] + jntid = dof_jntid[dofid] - jntid = dof_jntid[dofid] + # actuator-level gravity compensation, skip if added as passive force + if ngravcomp and jnt_actgravcomp[jntid]: + qfrc += qfrc_gravcomp_in[worldid, dofid] - # actuator-level gravity compensation, skip if added as passive force - if ngravcomp and jnt_actgravcomp[jntid]: - qfrc += qfrc_gravcomp_in[worldid, dofid] + if jnt_actfrclimited[jntid]: + frcrange = jnt_actfrcrange[worldid, jntid] + qfrc = wp.clamp(qfrc, frcrange[0], frcrange[1]) - if jnt_actfrclimited[jntid]: - frcrange = jnt_actfrcrange[worldid, jntid] - qfrc = wp.clamp(qfrc, frcrange[0], frcrange[1]) - - qfrc_actuator_out[worldid, dofid] = qfrc - - wp.launch_tiled( - qfrc_actuator, - dim=(d.nworld, m.nv), - inputs=[ - m.ngravcomp, - m.jnt_actfrclimited, - m.jnt_actfrcrange, - m.jnt_actgravcomp, - m.dof_jntid, - d.actuator_moment, - d.qfrc_gravcomp, - d.actuator_force, - ], - outputs=[d.qfrc_actuator], - block_dim=m.block_dim.qfrc_actuator, - ) + qfrc_actuator_out[worldid, dofid] = qfrc @event_scope @@ -923,7 +901,22 @@ def fwd_actuation(m: Model, d: Data): outputs=[d.actuator_force], ) - _qfrc_actuator(m, d) + wp.launch( + _qfrc_actuator, + dim=(d.nworld, m.nv), + inputs=[ + m.nu, + m.ngravcomp, + m.jnt_actfrclimited, + m.jnt_actfrcrange, + m.jnt_actgravcomp, + m.dof_jntid, + d.actuator_moment, + d.qfrc_gravcomp, + d.actuator_force, + ], + outputs=[d.qfrc_actuator], + ) @wp.kernel diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index 2441f8db..a6dafbf5 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -417,6 +417,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: nsite=mjm.nsite, ncam=mjm.ncam, nlight=mjm.nlight, + nmat=mjm.nmat, nflex=mjm.nflex, nflexvert=mjm.nflexvert, nflexedge=mjm.nflexedge, @@ -593,6 +594,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: light_mode=wp.array(mjm.light_mode, dtype=int), light_bodyid=wp.array(mjm.light_bodyid, dtype=int), light_targetbodyid=wp.array(mjm.light_targetbodyid, dtype=int), + light_type=create_nmodel_batched_array(mjm.light_type, dtype=int), + light_castshadow=create_nmodel_batched_array(mjm.light_castshadow, dtype=bool), + light_active=create_nmodel_batched_array(mjm.light_active, dtype=bool), light_pos=create_nmodel_batched_array(mjm.light_pos, dtype=wp.vec3), light_dir=create_nmodel_batched_array(mjm.light_dir, dtype=wp.vec3), light_poscom0=create_nmodel_batched_array(mjm.light_poscom0, dtype=wp.vec3), @@ -818,6 +822,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: plugin=wp.array(plugin_id, dtype=int), plugin_attr=wp.array(plugin_attr, dtype=wp.vec3f), geom_plugin_index=wp.array(geom_plugin_index, dtype=int), + mat_texid=create_nmodel_batched_array(mjm.mat_texid, dtype=int), + mat_texrepeat=create_nmodel_batched_array(mjm.mat_texrepeat, dtype=wp.vec2), mat_rgba=create_nmodel_batched_array(mjm.mat_rgba, dtype=wp.vec4), actuator_trntype_body_adr=wp.array(np.nonzero(mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY)[0], dtype=int), geompair2hfgeompair=wp.array(_hfield_geom_pair(mjm)[1], dtype=int), diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py index 7ebd0108..6ed403ef 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -44,7 +44,6 @@ class BlockDim: euler_dense: int = 256 actuator_velocity: int = 32 tendon_velocity: int = 256 - qfrc_actuator: int = 256 # ray ray: int = 64 # sensor @@ -724,6 +723,7 @@ class Model: nsite: number of sites ncam: number of cameras nlight: number of lights + nmat: number of materials nexclude: number of excluded geom pairs neq: number of equality constraints nmocap: number of mocap bodies @@ -856,6 +856,9 @@ class Model: light_mode: light tracking mode (CamLightType) (nlight,) light_bodyid: id of light's body (nlight,) light_targetbodyid: id of targeted body; -1: none (nlight,) + light_type: spot, directional, etc. (mjtLightType) (nworld, nlight) + light_castshadow: does light cast shadows (nworld, nlight) + light_active: is light active (nworld, nlight) light_pos: position rel. to body frame (nworld, nlight, 3) light_dir: direction rel. to body frame (nworld, nlight, 3) light_poscom0: global position rel. to sub-com in qpos0 (nworld, nlight, 3) @@ -864,11 +867,10 @@ class Model: mesh_vertadr: first vertex address (nmesh,) mesh_vertnum: number of vertices (nmesh,) mesh_vert: vertex positions for all meshes (nmeshvert, 3) - mesh_normal: normals for all meshes (nmeshnormal, 3) mesh_faceadr: first face address (nmesh,) mesh_face: face indices for all meshes (nface, 3) mesh_normaladr: first normal address (nmesh,) - mesh_normal: normals for all meshes (nmeshnormal x 3) + mesh_normal: normals for all meshes (nmeshnormal, 3) mesh_graphadr: graph data address; -1: no graph (nmesh,) mesh_graph: convex graph data (nmeshgraph,) mesh_quat: rotation applied to asset vertices (nmesh, 4) @@ -1001,6 +1003,8 @@ class Model: plugin_attr: config attributes of geom plugin (nplugin, 3) geom_plugin_index: geom index in plugin array (ngeom, ) mocap_bodyid: id of body for mocap (nmocap,) + mat_texid: texture id for rendering (nworld, nmat, mjNTEXROLE) + mat_texrepeat: texture repeat for rendering (nworld, nmat, 2) mat_rgba: rgba (nworld, nmat, 4) actuator_trntype_body_adr: addresses for actuators (<=nu,) with body transmission @@ -1021,6 +1025,7 @@ class Model: nsite: int ncam: int nlight: int + nmat: int nflex: int nflexvert: int nflexedge: int @@ -1160,6 +1165,9 @@ class Model: light_mode: wp.array(dtype=int) light_bodyid: wp.array(dtype=int) light_targetbodyid: wp.array(dtype=int) + light_type: wp.array2d(dtype=int) + light_castshadow: wp.array2d(dtype=bool) + light_active: wp.array2d(dtype=bool) light_pos: wp.array2d(dtype=wp.vec3) light_dir: wp.array2d(dtype=wp.vec3) light_poscom0: wp.array2d(dtype=wp.vec3) @@ -1312,6 +1320,8 @@ class Model: plugin_attr: wp.array(dtype=wp.vec3f) geom_plugin_index: wp.array(dtype=int) # warp only mocap_bodyid: wp.array(dtype=int) # warp only + mat_texid: wp.array3d(dtype=int) + mat_texrepeat: wp.array2d(dtype=wp.vec2) mat_rgba: wp.array2d(dtype=wp.vec4) actuator_trntype_body_adr: wp.array(dtype=int) # warp only geompair2hfgeompair: wp.array(dtype=int) # warp only diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index ca4f200a..0b90c028 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -42,7 +42,6 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) - @ffi.format_args_for_warp def _collision_shim( # Model diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 9aab9aba..d80eae30 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -42,7 +42,6 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) - @ffi.format_args_for_warp def _forward_shim( # Model @@ -1990,7 +1989,6 @@ _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 diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 3cb5f035..ea56b2b5 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -65,7 +65,6 @@ class BlockDim: mul_m_dense: int qderiv_actuator_passive_actuation: int qderiv_actuator_passive_no_actuation: int - qfrc_actuator: int ray: int segmented_sort: int tendon_velocity: int @@ -139,9 +138,11 @@ class ModelWarp(PyTreeNode): has_sdf_geom: bool jnt_limited_ball_adr: np.ndarray jnt_limited_slide_hinge_adr: np.ndarray + light_active: jax.Array light_bodyid: np.ndarray light_targetbodyid: np.ndarray mapM2M: np.ndarray + mat_texrepeat: jax.Array mesh_polyadr: np.ndarray mesh_polymap: np.ndarray mesh_polymapadr: np.ndarray @@ -648,7 +649,6 @@ _NDIM = { 'block_dim__mul_m_dense': 0, 'block_dim__qderiv_actuator_passive_actuation': 0, 'block_dim__qderiv_actuator_passive_no_actuation': 0, - 'block_dim__qfrc_actuator': 0, 'block_dim__ray': 0, 'block_dim__segmented_sort': 0, 'block_dim__tendon_velocity': 0, @@ -774,7 +774,9 @@ _NDIM = { 'jnt_solref': 3, 'jnt_stiffness': 2, 'jnt_type': 1, + 'light_active': 2, 'light_bodyid': 1, + 'light_castshadow': 2, 'light_dir': 3, 'light_dir0': 3, 'light_mode': 1, @@ -782,8 +784,11 @@ _NDIM = { 'light_pos0': 3, 'light_poscom0': 3, 'light_targetbodyid': 1, + 'light_type': 2, 'mapM2M': 1, 'mat_rgba': 3, + 'mat_texid': 3, + 'mat_texrepeat': 3, 'mesh_face': 2, 'mesh_faceadr': 1, 'mesh_graph': 1, @@ -823,6 +828,7 @@ _NDIM = { 'njnt': 0, 'nlight': 0, 'nlsp': 0, + 'nmat': 0, 'nmeshface': 0, 'nmeshgraph': 0, 'nmeshpoly': 0, @@ -1224,7 +1230,6 @@ _BATCH_DIM = { 'block_dim__mul_m_dense': False, 'block_dim__qderiv_actuator_passive_actuation': False, 'block_dim__qderiv_actuator_passive_no_actuation': False, - 'block_dim__qfrc_actuator': False, 'block_dim__ray': False, 'block_dim__segmented_sort': False, 'block_dim__tendon_velocity': False, @@ -1350,7 +1355,9 @@ _BATCH_DIM = { 'jnt_solref': True, 'jnt_stiffness': True, 'jnt_type': False, + 'light_active': True, 'light_bodyid': False, + 'light_castshadow': True, 'light_dir': True, 'light_dir0': True, 'light_mode': False, @@ -1358,8 +1365,11 @@ _BATCH_DIM = { 'light_pos0': True, 'light_poscom0': True, 'light_targetbodyid': False, + 'light_type': True, 'mapM2M': False, 'mat_rgba': True, + 'mat_texid': True, + 'mat_texrepeat': True, 'mesh_face': False, 'mesh_faceadr': False, 'mesh_graph': False, @@ -1399,6 +1409,7 @@ _BATCH_DIM = { 'njnt': False, 'nlight': False, 'nlsp': False, + 'nmat': False, 'nmeshface': False, 'nmeshgraph': False, 'nmeshpoly': False,