[BUG FIX] Various minor bug fixes. (#1368)

* Automatically kill irrelevant CI jobs.
* Fix partially broken taichi stdout silencing.
* Disable static_print in kernels because it does not work as intended.
* More robust inverse kinematics unit test.
* Get around MacOS Metal GPU taichi compiler bug (again!).
* Fix taichi type cast warnings.

Author: duburcqa

.github/workflows/generic-cpu.yml

@@ -5,6 +5,10 @@ on:
     branches:
       - main
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   generic-cpu:
     strategy:

.github/workflows/linux-gpu.yml

@@ -7,6 +7,12 @@ on:
       - main
   pull_request:
 
+concurrency:
+  # Cancel all workflows that are stil running if any when updating branches associated with PRs,
+  # BUT don't do anything for workflows that are not triggered by PRs.
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   linux-gpu:
     runs-on: [self-hosted, coreweave]

genesis/__init__.py

@@ -1,23 +1,19 @@
-# import taichi while suppressing its output
+import io
 import os
 import sys
 import site
 import atexit
 import logging as _logging
 import traceback
 from platform import system
-from unittest.mock import patch
+from contextlib import redirect_stdout
 
+# Import taichi while collecting its output without printing directly
+_ti_outputs = io.StringIO()
 
-_ti_outputs = []
+os.environ.setdefault("TI_ENABLE_PYBUF", "0" if sys.stdout is sys.__stdout__ else "1")
 
-
-def fake_print(*args, **kwargs):
-    output = "".join(args)
-    _ti_outputs.append(output)
-
-
-with patch("builtins.print", fake_print):
+with redirect_stdout(_ti_outputs):
     import taichi as ti
 
 try:
@@ -52,11 +48,9 @@ def init(
     logger_verbose_time=False,
     performance_mode: bool = False,  # True: compilation up to 6x slower (GJK), but runs ~1-5% faster
 ):
-    # Consider Genesis as initialized right away
     global _initialized
     if _initialized:
         raise_exception("Genesis already initialized.")
-    _initialized = True
 
     # genesis._theme
     global _theme
@@ -193,7 +187,7 @@ def init(
         ti_arch = TI_ARCH[platform][gs_backend.cpu]
 
     # init taichi
-    with patch("builtins.print", fake_print):
+    with redirect_stdout(_ti_outputs):
         ti.init(
             arch=ti_arch,
             # debug is causing segfault on some machines
@@ -227,9 +221,10 @@ def init(
         f"Running on ~~<[{device_name}]>~~ with backend ~~<{backend}>~~. Device memory: ~~<{total_mem:.2f}>~~ GB."
     )
 
-    for ti_output in _ti_outputs:
+    for ti_output in _ti_outputs.getvalue().splitlines():
         logger.debug(ti_output)
-    _ti_outputs.clear()
+    _ti_outputs.truncate(0)
+    _ti_outputs.seek(0)
 
     global exit_callbacks
     exit_callbacks = []
@@ -238,6 +233,8 @@ def init(
         f"🚀 Genesis initialized. 🔖 version: ~~<{__version__}>~~, 🌱 seed: ~~<{seed}>~~, 📏 precision: '~~<{precision}>~~', 🐛 debug: ~~<{debug}>~~, 🎨 theme: '~~<{theme}>~~'."
     )
 
+    _initialized = True
+
 
 ########################## init ##########################
 def destroy():

genesis/engine/bvh.py

@@ -343,38 +343,38 @@ def compute_bounds(self):
     @ti.kernel
     def _kernel_compute_bounds_init(self):
         self.updated[None] = True
-        self.internal_node_active.fill(0)
-        self.internal_node_ready.fill(0)
+        self.internal_node_active.fill(False)
+        self.internal_node_ready.fill(False)
 
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs):
             idx = ti.i32(self.morton_codes[i_b, i])
             self.nodes[i_b, i + self.n_aabbs - 1].bound.min = self.aabbs[i_b, idx].min
             self.nodes[i_b, i + self.n_aabbs - 1].bound.max = self.aabbs[i_b, idx].max
             parent_idx = self.nodes[i_b, i + self.n_aabbs - 1].parent
             if parent_idx != -1:
-                self.internal_node_active[i_b, parent_idx] = 1
+                self.internal_node_active[i_b, parent_idx] = True
 
     @ti.kernel
     def _kernel_compute_bounds_one_layer(self):
         self.updated[None] = False
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
-            if self.internal_node_active[i_b, i] == 0:
+            if not self.internal_node_active[i_b, i]:
                 continue
             left_bound = self.nodes[i_b, self.nodes[i_b, i].left].bound
             right_bound = self.nodes[i_b, self.nodes[i_b, i].right].bound
             self.nodes[i_b, i].bound.min = ti.min(left_bound.min, right_bound.min)
             self.nodes[i_b, i].bound.max = ti.max(left_bound.max, right_bound.max)
             parent_idx = self.nodes[i_b, i].parent
             if parent_idx != -1:
-                self.internal_node_ready[i_b, parent_idx] = 1
-            self.internal_node_active[i_b, i] = 0
+                self.internal_node_ready[i_b, parent_idx] = True
+            self.internal_node_active[i_b, i] = False
             self.updated[None] = True
 
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
-            if self.internal_node_ready[i_b, i] == 0:
+            if not self.internal_node_ready[i_b, i]:
                 continue
-            self.internal_node_active[i_b, i] = 1
-            self.internal_node_ready[i_b, i] = 0
+            self.internal_node_active[i_b, i] = True
+            self.internal_node_ready[i_b, i] = False
 
     @ti.kernel
     def query(self, aabbs: ti.template()):

genesis/engine/solvers/rigid/collider_decomp.py

@@ -153,6 +153,8 @@ def _init_collision_fields(self) -> None:
         self._max_collision_pairs = min(n_possible_pairs, self._solver._max_collision_pairs)
         self._max_contact_pairs = self._max_collision_pairs * self._n_contacts_per_pair
 
+        # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard
+        # condition, while 'print' is slowing down the kernel even if every called in practice...
         self._warn_msg_max_collision_pairs = (
             f"{colors.YELLOW}[Genesis] [00:00:00] [WARNING] Ignoring contact pair to avoid exceeding max "
             f"({self._max_contact_pairs}). Please increase the value of RigidSolver's option "
@@ -777,7 +779,7 @@ def _func_broad_phase(self):
                                 continue
 
                             if self.n_broad_pairs[i_b] == self._max_collision_pairs:
-                                ti.static_print(self._warn_msg_max_collision_pairs)
+                                # print(self._warn_msg_max_collision_pairs)
                                 break
                             self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][0] = i_ga
                             self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][1] = i_gb
@@ -887,22 +889,27 @@ def _func_narrow_phase_convex_vs_convex(self):
                 i_ga = self.broad_collision_pairs[i_pair, i_b][0]
                 i_gb = self.broad_collision_pairs[i_pair, i_b][1]
 
+                if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type:
+                    i_ga, i_gb = i_gb, i_ga
+
                 if (
                     self._solver.geoms_info[i_ga].is_convex
                     and self._solver.geoms_info[i_gb].is_convex
                     and not self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.TERRAIN
-                    and not (
-                        self._solver._enable_multi_contact
-                        and self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
-                        and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
-                    )
                     and not (
                         self._solver._box_box_detection
                         and self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.BOX
                         and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
                     )
                 ):
-                    self._func_convex_convex_contact(i_ga, i_gb, i_b)
+                    if ti.static(sys.platform == "darwin"):
+                        self._func_convex_convex_contact(i_ga, i_gb, i_b)
+                    else:
+                        if not (
+                            self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
+                            and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
+                        ):
+                            self._func_convex_convex_contact(i_ga, i_gb, i_b)
 
     @ti.kernel
     def _func_narrow_phase_convex_specializations(self):
@@ -915,15 +922,11 @@ def _func_narrow_phase_convex_specializations(self):
                 if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type:
                     i_ga, i_gb = i_gb, i_ga
 
-                if (
-                    self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
-                    and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
-                ):
-                    if ti.static(sys.platform == "darwin"):
-                        # FIXME: It seems redundant, why don't we just call _func_plane_box_contact directly?
-                        # Anyway in this function, we will call _func_plane_box_contact.
-                        self._func_convex_convex_contact(i_ga, i_gb, i_b)
-                    else:
+                if ti.static(sys.platform != "darwin"):
+                    if (
+                        self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
+                        and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
+                    ):
                         self._func_plane_box_contact(i_ga, i_gb, i_b)
 
                 if ti.static(self._solver._box_box_detection):
@@ -1110,7 +1113,10 @@ def _func_add_contact(self, i_ga, i_gb, normal, contact_pos, penetration, i_b):
         i_col = self.n_contacts[i_b]
 
         if i_col == self._max_contact_pairs:
-            ti.static_print(self._warn_msg_max_collision_pairs)
+            # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard
+            # condition, while 'print' is slowing down the kernel even if every called in practice...
+            # print(self._warn_msg_max_collision_pairs)
+            pass
         else:
             ga_info = self._solver.geoms_info[i_ga]
             gb_info = self._solver.geoms_info[i_gb]
@@ -1194,9 +1200,6 @@ def _func_contact_orthogonals(self, i_ga, i_gb, normal, i_b):
 
     @ti.func
     def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
-        if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type:
-            i_ga, i_gb = i_gb, i_ga
-
         if (
             self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
             and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX

tests/test_bvh.py

@@ -1,11 +1,10 @@
-import taichi as ti
 import torch
 import numpy as np
-import genesis as gs
+import pytest
 
+import genesis as gs
 from genesis.engine.bvh import LBVH, AABB
 
-import pytest
 from .utils import assert_allclose
 
 
@@ -37,18 +36,19 @@ def test_morton_code(lbvh):
             ), f"Morton codes are not sorted: {morton_codes[i_b, i]} < {morton_codes[i_b, i - 1]}"
 
 
-@ti.kernel
-def expand_bits(lbvh: ti.template(), x: ti.template(), expanded_x: ti.template()):
-    n_x = x.shape[0]
-    for i in range(n_x):
-        expanded_x[i] = lbvh.expand_bits(x[i])
-
-
 def test_expand_bits():
     """
     Test the expand_bits function for LBVH.
     A 10-bit integer is expanded to a 30-bit integer by inserting two zeros before each bit.
     """
+    import taichi as ti
+
+    @ti.kernel
+    def expand_bits(lbvh: ti.template(), x: ti.template(), expanded_x: ti.template()):
+        n_x = x.shape[0]
+        for i in range(n_x):
+            expanded_x[i] = lbvh.expand_bits(x[i])
+
     # random integer
     x_np = np.random.randint(0, 1024, (10,), dtype=np.uint32)
     x_ti = ti.field(ti.uint32, shape=x_np.shape)

tests/test_rigid_physics.py

@@ -14,7 +14,7 @@
 
 import genesis as gs
 import genesis.utils.geom as gu
-from genesis.utils.misc import tensor_to_array
+from genesis.utils.misc import tensor_to_array, get_assets_dir
 
 from .utils import (
     get_hf_assets,
@@ -1404,14 +1404,14 @@ def move_cube(use_suction, mode, show_viewer):
     # move to pre-grasp pose
     qpos = franka.inverse_kinematics(
         link=end_effector,
-        pos=np.array([0.65, 0.0, 0.25]),
+        pos=np.array([0.65, 0.0, 0.22]),
         quat=np.array([0, 1, 0, 0]),
     )
     # gripper open pos
     qpos[-2:] = 0.04
     path = franka.plan_path(
         qpos_goal=qpos,
-        num_waypoints=100,  # 1s duration
+        num_waypoints=120,  # 1s duration
         resolution=0.05,
         timeout=30.0,
     )
@@ -1423,7 +1423,7 @@ def move_cube(use_suction, mode, show_viewer):
         scene.step()
 
     # Get more time to the robot to reach the last waypoint
-    for i in range(100):
+    for i in range(120):
         scene.step()
 
     # reach
@@ -1433,7 +1433,7 @@ def move_cube(use_suction, mode, show_viewer):
         quat=np.array([0, 1, 0, 0]),
     )
     franka.control_dofs_position(qpos[:-2], motors_dof)
-    for i in range(50):
+    for i in range(60):
         scene.step()
 
     # grasp
@@ -1465,7 +1465,7 @@ def move_cube(use_suction, mode, show_viewer):
     )
     path = franka.plan_path(
         qpos_goal=qpos,
-        num_waypoints=150,
+        num_waypoints=80,
         resolution=0.05,
         timeout=30.0,
     )
@@ -1488,7 +1488,7 @@ def move_cube(use_suction, mode, show_viewer):
         scene.step()
         if i > 550:
             qvel = cube.get_dofs_velocity()
-            assert_allclose(qvel, 0, atol=0.06)
+            assert_allclose(qvel, 0, atol=0.02)
 
     qpos = cube.get_dofs_position()
     assert_allclose(qpos[2], 0.075, atol=2e-3)
@@ -2715,28 +2715,29 @@ def must_cast(value):
 
 
 @pytest.mark.parametrize("backend", [gs.cpu])
-def test_mesh_to_heightfield(show_viewer):
+def test_mesh_to_heightfield(tmp_path, show_viewer):
+    horizontal_scale = 2.0
+    path_terrain = os.path.join(get_assets_dir(), "meshes", "terrain_45.obj")
+
+    hf_terrain, xs, ys = gs.utils.terrain.mesh_to_heightfield(path_terrain, spacing=horizontal_scale, oversample=1)
+
+    # default heightfield starts at 0, 0, 0
+    # translate to the center of the mesh
+    translation = np.array([np.nanmin(xs), np.nanmin(ys), 0])
+
     ########################## create a scene ##########################
     scene = gs.Scene(
-        show_viewer=show_viewer,
         sim_options=gs.options.SimOptions(
             gravity=(2, 0, -2),
         ),
         viewer_options=gs.options.ViewerOptions(
             camera_pos=(0, -50, 0),
             camera_lookat=(0, 0, 0),
         ),
+        show_viewer=show_viewer,
+        show_FPS=False,
     )
 
-    horizontal_scale = 2.0
-    gs_root = os.path.dirname(os.path.abspath(gs.__file__))
-    path_terrain = os.path.join(gs_root, "assets", "meshes", "terrain_45.obj")
-    hf_terrain, xs, ys = gs.utils.terrain.mesh_to_heightfield(path_terrain, spacing=horizontal_scale, oversample=1)
-
-    # default heightfield starts at 0, 0, 0
-    # translate to the center of the mesh
-    translation = np.array([np.nanmin(xs), np.nanmin(ys), 0])
-
     terrain_heightfield = scene.add_entity(
         morph=gs.morphs.Terrain(
             horizontal_scale=horizontal_scale,
@@ -2746,15 +2747,13 @@ def test_mesh_to_heightfield(show_viewer):
         ),
         vis_mode="collision",
     )
-
     ball = scene.add_entity(
         gs.morphs.Sphere(
             pos=(10, 15, 10),
             radius=1,
         ),
         vis_mode="collision",
     )
-
     scene.build()
 
     for i in range(1000):