From 1feaf8fd672a2a07e69f61ca8d84d92f37c59041 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 15 Nov 2024 16:54:44 -0800
Subject: [PATCH] Add *rgba fields as jax.Arrays for visual domain rando.

PiperOrigin-RevId: 697038459
Change-Id: Ic6eae248e916329daa8fbbeb0770cf689f06b87b
---
 mjx/mujoco/mjx/_src/ray.py   | 5 +++--
 mjx/mujoco/mjx/_src/types.py | 4 ++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py
index 9335f4e9..6d000373 100644
--- a/mjx/mujoco/mjx/_src/ray.py
+++ b/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:
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index a940d44e..dcfbdb6f 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/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