From 1b93226fb046ae906387887a3fc9ce1e74bc229b Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Fri, 16 Feb 2024 12:39:30 -0800
Subject: [PATCH] To save memory, avoid loading meshes that will never be used
 for collisions.

PiperOrigin-RevId: 607770456
Change-Id: I34c4a34d1820a9fd88d2ebeaffec9a64637d2614
---
 mjx/mujoco/mjx/_src/mesh.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py
index eaa5d3bf..bebca9ac 100644
--- a/mjx/mujoco/mjx/_src/mesh.py
+++ b/mjx/mujoco/mjx/_src/mesh.py
@@ -212,10 +212,14 @@ def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
   """Derives geom mesh attributes for mjx.Model from MjModel."""
   kwargs = {k: [] for k in _DERIVED_ARGS}
   verts, faces = _get_faces_verts(m)
+  geom_con = m.geom_conaffinity | m.geom_contype
   for geomid in range(m.ngeom):
     dataid = m.geom_dataid[geomid]
-    typ = m.geom_type[geomid]
-    if typ == GeomType.BOX:
+    if not geom_con[geomid]:
+      # ignore visual-only meshes
+      kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
+      continue
+    elif m.geom_type[geomid] == GeomType.BOX:
       vert, face = _box(m.geom_size[geomid])
     elif dataid >= 0:
       vert, face = verts[dataid], faces[dataid]