Add *rgba fields as jax.Arrays for visual domain rando.

PiperOrigin-RevId: 697038459
Change-Id: Ic6eae248e916329daa8fbbeb0770cf689f06b87b

mjx/mujoco/mjx/_src/ray.py

@@ -257,8 +257,6 @@ def ray(
 
   dists, ids = [], []
   geom_filter = m.geom_bodyid != bodyexclude
-  geom_filter &= (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0)
-  geom_filter &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0)
   geom_filter &= flg_static | (m.body_weldid[m.geom_bodyid] != 0)
   if geomgroup:
     geomgroup = np.array(geomgroup, dtype=bool)
@@ -268,6 +266,8 @@ def ray(
   geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos)
   geom_vecs = jax.vmap(lambda x: x.T @ vec)(d.geom_xmat)
 
+  geom_filter_dyn = (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0)
+  geom_filter_dyn &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0)
   for geom_type, fn in _RAY_FUNC.items():
     id_, = np.nonzero(geom_filter & (m.geom_type == geom_type))
 
@@ -281,6 +281,7 @@ def ray(
     else:
       dist = jax.vmap(fn)(*args)
 
+    dist = jp.where(geom_filter_dyn[id_], dist, jp.inf)
     dists, ids = dists + [dist], ids + [id_]
 
   if not ids:

mjx/mujoco/mjx/_src/types.py

@@ -946,7 +946,7 @@ class Model(PyTreeNode):
   geom_margin: jax.Array
   geom_gap: jax.Array
   geom_fluid: np.ndarray
-  geom_rgba: np.ndarray
+  geom_rgba: jax.Array
   site_type: np.ndarray
   site_bodyid: np.ndarray
   site_sameframe: np.ndarray
@@ -1040,7 +1040,7 @@ class Model(PyTreeNode):
   tex_nchannel: np.ndarray
   tex_adr: np.ndarray
   tex_data: jax.Array
-  mat_rgba: np.ndarray
+  mat_rgba: jax.Array
   mat_texid: np.ndarray
   pair_dim: np.ndarray
   pair_geom1: np.ndarray