[BUG FIX] Fix URDF color overwrite. (#2065)

* Fix URDF color overwrite.
* Fix zero-copy being disabled for numpy.
* Add warning when enabling requires_grad while using GsTaichi dynamic arrays.
* More robust unit tests.

Author: duburcqa

genesis/engine/scene.py

@@ -276,6 +276,11 @@ def _validate_options(
             gs.raise_exception("`renderer` should be an instance of `gs.renderers.Renderer`.")
 
         # Validate rigid_options against sim_options
+        if sim_options.requires_grad and gs.use_ndarray:
+            gs.logger.warning(
+                "Use GsTaichi dynamic array mode while enabling gradient computation is not recommended. Please "
+                "enable performance mode at init for efficiency, e.g. 'gs.init(..., performance_mode=True)'."
+            )
         if rigid_options.box_box_detection is None:
             rigid_options.box_box_detection = not sim_options.requires_grad
         elif rigid_options.box_box_detection and sim_options.requires_grad:

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -258,7 +258,6 @@ def build(self):
                 sparse_solve=self._options.sparse_solve,
                 integrator=self._integrator,
                 solver_type=self._options.constraint_solver,
-                is_backward=False,
             )
             # Add terms for static inner loops, use -1 if not requires_grad to avoid re-compilation
             if self.sim.options.requires_grad:
@@ -281,7 +280,6 @@ def build(self):
                 enable_collision=self._enable_collision,
                 integrator=gs.integrator.approximate_implicitfast,
                 solver_type=gs.constraint_solver.CG,
-                is_backward=False,
             )
 
         if self._static_rigid_sim_config.requires_grad:
@@ -7103,11 +7101,9 @@ def kernel_update_geoms_render_T(
     ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
     for i_g, i_b in ti.ndrange(n_geoms, _B):
         geom_T = gu.ti_trans_quat_to_T(
-            geoms_state.pos[i_g, i_b] + rigid_global_info.envs_offset[i_b],
-            geoms_state.quat[i_g, i_b],
-            EPS,
+            geoms_state.pos[i_g, i_b] + rigid_global_info.envs_offset[i_b], geoms_state.quat[i_g, i_b], EPS
         )
-        for J in ti.static(ti.grouped(ti.static(ti.ndrange(4, 4)))):
+        for J in ti.static(ti.grouped(ti.ndrange(4, 4))):
             geoms_render_T[(i_g, i_b, *J)] = ti.cast(geom_T[J], ti.float32)
 
 
@@ -7127,9 +7123,7 @@ def kernel_update_vgeoms_render_T(
     ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
     for i_g, i_b in ti.ndrange(n_vgeoms, _B):
         geom_T = gu.ti_trans_quat_to_T(
-            vgeoms_state.pos[i_g, i_b] + rigid_global_info.envs_offset[i_b],
-            vgeoms_state.quat[i_g, i_b],
-            EPS,
+            vgeoms_state.pos[i_g, i_b] + rigid_global_info.envs_offset[i_b], vgeoms_state.quat[i_g, i_b], EPS
         )
         for J in ti.static(ti.grouped(ti.ndrange(4, 4))):
             vgeoms_render_T[(i_g, i_b, *J)] = ti.cast(geom_T[J], ti.float32)

genesis/utils/array_class.py

@@ -1795,10 +1795,10 @@ class StructRigidSimStaticConfig(metaclass=AutoInitMeta):
     batch_links_info: bool = False
     batch_dofs_info: bool = False
     batch_joints_info: bool = False
-    enable_mujoco_compatibility: bool = False
+    enable_mujoco_compatibility: bool = True
     enable_multi_contact: bool = False
     enable_joint_limit: bool = False
-    box_box_detection: bool = True
+    box_box_detection: bool = False
     sparse_solve: bool = False
     integrator: int = gs.integrator.approximate_implicitfast
     solver_type: int = gs.constraint_solver.CG

genesis/utils/misc.py

@@ -553,11 +553,10 @@ def ti_to_python(
     data_type = type(value)
     use_zerocopy = (
         gs.use_zerocopy
-        and (to_torch or gs.backend == gs.cpu)
         and not issubclass(data_type, ti.Field)
         # and (gs.backend != gs.metal or not issubclass(data_type, ti.Field))
     )
-    if not use_zerocopy:
+    if not use_zerocopy or (not to_torch and gs.backend != gs.cpu):
         if copy is False:
             gs.raise_exception(
                 "Specifying 'copy=False' is not supported if 'gs.use_zerocopy=False' or ('to_torch=False' and "
@@ -570,21 +569,22 @@ def ti_to_python(
     # Leverage zero-copy if enabled
     batch_shape = value.shape
     if use_zerocopy:
-        try:
-            if to_torch or gs.backend != gs.cpu:
-                out = value._T_tc if transpose else value._tc
-            else:
-                out = value._T_np if transpose else value._np
-        except AttributeError:
-            value._tc = torch.utils.dlpack.from_dlpack(value.to_dlpack())
-            value._T_tc = value._tc.movedim(batch_ndim - 1, 0) if (batch_ndim := len(batch_shape)) > 1 else value._tc
-            if to_torch:
-                out = value._T_tc if transpose else value._tc
-            if gs.backend == gs.cpu:
-                value._np = value._tc.numpy()
-                value._T_np = value._T_tc.numpy()
-                if not to_torch:
+        while True:
+            try:
+                if to_torch or gs.backend != gs.cpu:
+                    out = value._T_tc if transpose else value._tc
+                else:
                     out = value._T_np if transpose else value._np
+                break
+            except AttributeError:
+                value._tc = torch.utils.dlpack.from_dlpack(value.to_dlpack())
+                if (batch_ndim := len(batch_shape)) > 1:
+                    value._T_tc = value._tc.movedim(batch_ndim - 1, 0)
+                else:
+                    value._T_tc = value._tc
+                if gs.backend == gs.cpu:
+                    value._np = value._tc.numpy()
+                    value._T_np = value._T_tc.numpy()
         if copy:
             if to_torch:
                 out = out.clone()

genesis/utils/urdf.py

@@ -2,8 +2,10 @@
 from itertools import chain
 from pathlib import Path
 
-import trimesh
 import numpy as np
+import trimesh
+from trimesh.visual import ColorVisuals, TextureVisuals
+from trimesh.visual.color import VertexColor
 
 import genesis as gs
 from genesis.ext import urdfpy
@@ -131,7 +133,11 @@ def parse_urdf(morph, surface):
                                 "'parse_glb_with_zup=True' in morph options if you find the mesh is 90-degree rotated. "
                             )
 
-                    if not geom_is_col and (morph.prioritize_urdf_material or not tmesh.visual.defined):
+                    visual = mesh.trimesh.visual
+                    has_color_override = (isinstance(visual, (ColorVisuals, TextureVisuals)) and visual.defined) or (
+                        isinstance(visual, VertexColor) and visual.vertex_colors.size > 0
+                    )
+                    if not geom_is_col and (morph.prioritize_urdf_material or not has_color_override):
                         if geom.material is not None and geom.material.color is not None:
                             mesh.set_color(geom.material.color)
 

tests/test_examples.py

@@ -24,6 +24,7 @@
 IGNORE_SCRIPT_NAMES = {
     "ddp_multi_gpu.py",
     "multi_gpu.py",
+    "visualization.py",  # FIXME: Flaky because of possibly undefined shadow map when running in thread
     "single_franka_batch_render.py",  # FIXME: segfault on exit
     "fem_cube_linked_with_arm.py",  # FIXME: segfault on exit (corrupted double-linked list)
 }
@@ -32,7 +33,7 @@
         "cut_dragon.py",
     }
 
-TIMEOUT = 450.0
+TIMEOUT = 500.0
 
 
 pytestmark = [

tests/test_grad.py

@@ -407,6 +407,7 @@ def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force):
         assert_allclose(dL_error, 0.0, atol=RTOL)
 
 
+@pytest.mark.slow  # ~250s
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu, gs.gpu])
 def test_differentiable_rigid(show_viewer):

tests/test_pbd.py

@@ -29,13 +29,22 @@ def test_maxvolume(pbd_material, show_viewer, box_obj_path):
 
     # Mesh without any maximum-element-volume constraint
     pbd1 = scene.add_entity(
-        morph=gs.morphs.Mesh(file=box_obj_path, nobisect=False, verbose=1),
+        morph=gs.morphs.Mesh(
+            file=box_obj_path,
+            nobisect=False,
+            verbose=1,
+        ),
         material=pbd_material,
     )
 
     # Mesh with maximum element volume limited to 0.001
     pbd2 = scene.add_entity(
-        morph=gs.morphs.Mesh(file=box_obj_path, nobisect=False, maxvolume=0.001, verbose=1),
+        morph=gs.morphs.Mesh(
+            file=box_obj_path,
+            nobisect=False,
+            maxvolume=0.001,
+            verbose=1,
+        ),
         material=pbd_material,
     )
 
@@ -154,18 +163,37 @@ def test_cloth_attach_rigid_link(show_viewer):
     box_height = 2.25
 
     scene = gs.Scene(
-        sim_options=gs.options.SimOptions(dt=2e-2, substeps=10, gravity=(0.0, 0.0, 0.0)),
-        pbd_options=gs.options.PBDOptions(particle_size=particle_size),
+        sim_options=gs.options.SimOptions(
+            dt=2e-2,
+            substeps=10,
+            gravity=(0.0, 0.0, 0.0),
+        ),
+        pbd_options=gs.options.PBDOptions(
+            particle_size=particle_size,
+        ),
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(2.5, 2.0, 4.0),
+            camera_lookat=(-1.0, -1.0, 1.0),
+        ),
         show_viewer=show_viewer,
     )
-    scene.add_entity(gs.morphs.Plane(), gs.materials.Rigid(needs_coup=True, coup_friction=0.0))
+    scene.add_entity(
+        gs.morphs.Plane(),
+        material=gs.materials.Rigid(
+            needs_coup=True,
+            coup_friction=0.0,
+        ),
+    )
 
     box = scene.add_entity(
         morph=gs.morphs.Box(
             pos=(0.25, 0.25, box_height),
             size=(0.25, 0.25, 0.25),
         ),
-        material=gs.materials.Rigid(needs_coup=True, coup_friction=0.0),
+        material=gs.materials.Rigid(
+            needs_coup=True,
+            coup_friction=0.0,
+        ),
     )
     cloth = scene.add_entity(
         morph=gs.morphs.Mesh(
@@ -174,9 +202,10 @@ def test_cloth_attach_rigid_link(show_viewer):
             scale=0.5,
         ),
         material=gs.materials.PBD.Cloth(),
-        surface=gs.surfaces.Default(color=(0.2, 0.4, 0.8, 1.0)),
+        surface=gs.surfaces.Default(
+            color=(0.2, 0.4, 0.8, 1.0),
+        ),
     )
-
     scene.build(n_envs=2)
 
     particles_idx = [0, 1, 2, 3, 4, 5, 6, 7]
@@ -190,38 +219,28 @@ def test_cloth_attach_rigid_link(show_viewer):
     cloth_pos0 = cloth.get_particles_pos()[:, particles_idx].clone()
     link_pos0 = scene.rigid_solver.links[box_link_idx].get_pos().clone()
 
-    for _ in range(25):
+    for _ in range(20):
         scene.step()
 
     # wait for the box to stop
     box.set_dofs_velocity(np.zeros_like(vel), dofs_idx_local=[0, 1, 2])
-    for _ in range(5):
+    for _ in range(10):
         scene.step()
 
     # Check that the attached particles followed the link displacement per env
     cloth_pos1 = cloth.get_particles_pos()[:, particles_idx].clone()
     link_pos1 = scene.rigid_solver.links[box_link_idx].get_pos().clone()
+    assert_allclose((cloth_pos1 - cloth_pos0).movedim(0, -2), link_pos1 - link_pos0, atol=2e-5)
 
-    cloth_disp = cloth_pos1 - cloth_pos0
-    link_disp = link_pos1 - link_pos0
-    # broadcast link_disp to match cloth_disp shape
-    link_disp = link_disp.unsqueeze(1)
-
-    assert_allclose(cloth_disp, link_disp, atol=2e-5)
-
-    # Release cloth and revert box's speed
+    # Release cloth and restore box's speed
     box.set_dofs_velocity(vel, dofs_idx_local=[0, 1, 2])
     cloth.release_particle(particles_idx)
-    for i in range(25):
+    for _ in range(15):
         scene.step()
 
     # Make sure that the cloth is laying on the ground without moving
-    cloth_pos2 = cloth.get_particles_pos()[:, particles_idx].clone()
-    link_pos2 = scene.rigid_solver.links[box_link_idx].get_pos().clone()
-    cloth_disp = cloth_pos2 - cloth_pos1
+    cloth_pos2 = cloth.get_particles_pos()[:, particles_idx]
+    link_pos2 = scene.rigid_solver.links[box_link_idx].get_pos()
     link_disp = link_pos2 - link_pos1
-    link_disp = link_disp.unsqueeze(1)
-
-    link_disp = link_disp.norm(dim=-1)
-    cloth_disp = cloth_disp.norm(dim=-1)
-    assert (link_disp - cloth_disp).norm(dim=-1).mean() > 0.1
+    cloth_disp = cloth_pos2 - cloth_pos1
+    assert ((cloth_disp.movedim(0, -2) - link_disp).norm(dim=-1) > 0.2).all()

tests/test_render.py

@@ -637,9 +637,13 @@ def test_camera_follow_entity(n_envs, renderer, show_viewer):
             GUI=show_viewer,
         )
         cam.follow_entity(obj, smoothing=None)
+        cam.unfollow_entity()
 
     scene.build(n_envs=n_envs)
 
+    for cam, obj in zip(scene.visualizer.cameras, scene.entities):
+        cam.follow_entity(obj, smoothing=None)
+
     # First render
     seg_mask = None
     for entity_idx, cam in enumerate(scene.visualizer.cameras, 1):