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

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 800210640
Change-Id: I06d05ee2560f54d975458f0e13b37157c088c33b
This commit is contained in:
Baruch Tabanpour
2025-08-27 16:10:13 -07:00
committed by Copybara-Service
parent 79f2bbc727
commit ba47402853
6 changed files with 77 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+44 -51
View File
@@ -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
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+6
View File
@@ -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),
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+13 -3
View File
@@ -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
mjx/mujoco/mjx/warp/collision_driver.py
-1
View File
@@ -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
mjx/mujoco/mjx/warp/forward.py
-2
View File
@@ -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
mjx/mujoco/mjx/warp/types.py
+14 -3
View File
@@ -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,