Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 800210640
Change-Id: I06d05ee2560f54d975458f0e13b37157c088c33b

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):
+@wp.kernel
-  NU = m.nu
+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
+  qfrc = float(0.0)
-  def qfrc_actuator(
+  for uid in range(nu):
-    # Model:
+    qfrc += actuator_moment_in[worldid, uid, dofid] * actuator_force_in[worldid, uid]
-    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()
 
-    actuator_moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, 1), offset=(0, dofid))
+  jntid = dof_jntid[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]
+  # 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 jnt_actfrclimited[jntid]:
-    if ngravcomp and jnt_actgravcomp[jntid]:
+    frcrange = jnt_actfrcrange[worldid, jntid]
-      qfrc += qfrc_gravcomp_in[worldid, dofid]
+    qfrc = wp.clamp(qfrc, frcrange[0], frcrange[1])
 
-    if jnt_actfrclimited[jntid]:
+  qfrc_actuator_out[worldid, dofid] = qfrc
-      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,
-  )
 
 
 @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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,