From f2e9e38edfc05f25fe93735ca8eb6aa16494fbdf Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Wed, 27 Nov 2024 08:12:35 -0800
Subject: [PATCH] Set geom_matid as a jax.Array for visual domain
 randomization.

PiperOrigin-RevId: 700706160
Change-Id: I13ef4da8cd3a671446e4bdd186dcacb4d0b3fb00
---
 mjx/mujoco/mjx/_src/types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 89dab88e..6446cd33 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -933,7 +933,7 @@ class Model(PyTreeNode):
   geom_sameframe: np.ndarray
   geom_dataid: np.ndarray
   geom_group: np.ndarray
-  geom_matid: np.ndarray
+  geom_matid: jax.Array
   geom_priority: np.ndarray
   geom_solmix: jax.Array
   geom_solref: jax.Array